robot viz in genjax (GEN-884) (#20)

Author: mhuebert

poetry.lock

@@ -12,9 +12,8 @@ packages = [
 [tool.poetry.dependencies]
 python = ">=3.11,<3.13"
 jupytext = "^1.16.1"
-genjax = {version = "0.7.0.post4.dev0+eacb241e",   source = "gcp" }
-# genstudio = {version = "2024.12.003", source = "gcp"}
-genstudio = {path = "../genstudio", develop = true}
+genjax = {version = "0.8.0",   source = "gcp" }
+genstudio = {version = "2024.12.004", source = "gcp"}
 ipykernel = "^6.29.3"
 matplotlib = "^3.8.3"
 anywidget = "^0.9.7"

pyproject.toml

@@ -0,0 +1,290 @@
+# %%
+import jax.numpy as jnp
+import time
+from robot_2.where_am_i import (
+    RobotCapabilities,
+    simulate_robot,
+    World,
+    walls_to_jax,
+    PRNGKey,
+)
+import jax.random
+
+# Benchmark State:
+# Walls:
+walls = [
+    [0, 0, 0],
+    [10, 0, 0],
+    [10, 0, 0],
+    [10, 10, 0],
+    [10, 10, 0],
+    [0, 10, 0],
+    [0, 10, 0],
+    [0, 0, 0],
+    [8.125, 9.03125, 1974574552],
+    [7.875, 9.0625, 1974574552],
+    [7.53125, 9.0625, 1974574552],
+    [7.09375, 9.09375, 1974574552],
+    [6.8125, 9.09375, 1974574552],
+    [6.5, 9.125, 1974574552],
+    [6.1875, 9.125, 1974574552],
+    [5.71875, 9.15625, 1974574552],
+    [5.40625, 9.15625, 1974574552],
+    [5.09375, 9.15625, 1974574552],
+    [4.75, 9.15625, 1974574552],
+    [4.4375, 9.125, 1974574552],
+    [4, 9.09375, 1974574552],
+    [3.71875, 9.0625, 1974574552],
+    [3.46875, 9.03125, 1974574552],
+    [3.21875, 9, 1974574552],
+    [2.96875, 8.9375, 1974574552],
+    [2.5625, 8.84375, 1974574552],
+    [2.09375, 8.75, 1974574552],
+    [1.71875, 8.625, 1974574552],
+    [1.75, 8.34375, 1974574552],
+    [2.03125, 8.25, 1974574552],
+    [2.34375, 8.1875, 1974574552],
+    [2.71875, 8.15625, 1974574552],
+    [3.0625, 8.125, 1974574552],
+    [3.40625, 8.09375, 1974574552],
+    [3.96875, 8.0625, 1974574552],
+    [4.3125, 8.03125, 1974574552],
+    [4.625, 8.03125, 1974574552],
+    [4.90625, 8, 1974574552],
+    [5.3125, 8, 1974574552],
+    [5.59375, 8, 1974574552],
+    [5.875, 7.9375, 1974574552],
+    [6.15625, 7.875, 1974574552],
+    [6.4375, 7.75, 1974574552],
+    [6.6875, 7.59375, 1974574552],
+    [6.8125, 7.375, 1974574552],
+    [6.6875, 7.15625, 1974574552],
+    [6.40625, 7, 1974574552],
+    [5.78125, 6.875, 1974574552],
+    [5.4375, 6.8125, 1974574552],
+    [5.09375, 6.75, 1974574552],
+    [4.34375, 6.625, 1974574552],
+    [3.90625, 6.5625, 1974574552],
+    [3.65625, 6.5, 1974574552],
+    [3.1875, 6.4375, 1974574552],
+    [2.625, 6.28125, 1974574552],
+    [2.65625, 5.96875, 1974574552],
+    [2.9375, 5.84375, 1974574552],
+    [3.25, 5.78125, 1974574552],
+    [3.53125, 5.75, 1974574552],
+    [4.03125, 5.71875, 1974574552],
+    [4.46875, 5.6875, 1974574552],
+    [5.21875, 5.6875, 1974574552],
+    [5.75, 5.6875, 1974574552],
+    [6.0625, 5.6875, 1974574552],
+    [6.3125, 5.6875, 1974574552],
+    [6.5625, 5.65625, 1974574552],
+    [6.59375, 5.3125, 1974574552],
+    [6.375, 5.1875, 1974574552],
+    [6, 5.0625, 1974574552],
+    [5.3125, 4.9375, 1974574552],
+    [4.9375, 4.875, 1974574552],
+    [4.1875, 4.8125, 1974574552],
+    [3.875, 4.75, 1974574552],
+    [3.46875, 4.71875, 1974574552],
+    [3.0625, 4.65625, 1974574552],
+    [2.75, 4.59375, 1974574552],
+    [2.46875, 4.5, 1974574552],
+    [2.78125, 4.1875, 1974574552],
+    [3.15625, 4.125, 1974574552],
+    [3.53125, 4.0625, 1974574552],
+    [4, 4, 1974574552],
+    [4.46875, 3.9375000000000004, 1974574552],
+    [5.25, 3.8749999999999996, 1974574552],
+    [5.75, 3.84375, 1974574552],
+    [6, 3.84375, 1974574552],
+    [6.28125, 3.84375, 1974574552],
+    [6.75, 3.84375, 1974574552],
+    [6.5, 3.53125, 1974574552],
+    [6.25, 3.5, 1974574552],
+    [5.875, 3.4375, 1974574552],
+    [5.59375, 3.4062499999999996, 1974574552],
+    [5.0625, 3.375, 1974574552],
+    [4.625, 3.375, 1974574552],
+    [3.9375, 3.34375, 1974574552],
+    [2.875, 3.3125000000000004, 1974574552],
+    [1.75, 3.1562500000000004, 1974574552],
+    [1.5, 3.0625, 1974574552],
+    [1.53125, 2.7812499999999996, 1974574552],
+    [1.8125, 2.65625, 1974574552],
+    [2.0625, 2.6249999999999996, 1974574552],
+    [2.625, 2.5, 1974574552],
+    [2.875, 2.4687499999999996, 1974574552],
+    [3.375, 2.40625, 1974574552],
+    [3.9375, 2.3750000000000004, 1974574552],
+    [4.8125, 2.34375, 1974574552],
+    [5.34375, 2.34375, 1974574552],
+    [5.875, 2.34375, 1974574552],
+    [6.34375, 2.34375, 1974574552],
+    [6.59375, 2.34375, 1974574552],
+    [7.15625, 2.34375, 1974574552],
+    [7.4375, 2.28125, 1974574552],
+    [7.1875, 2.03125, 1974574552],
+    [6.75, 1.9374999999999998, 1974574552],
+    [6.4375, 1.875, 1974574552],
+    [5.59375, 1.7812499999999998, 1974574552],
+    [5.125, 1.7812499999999998, 1974574552],
+    [4.625, 1.7812499999999998, 1974574552],
+    [3.96875, 1.71875, 1974574552],
+    [3.4375, 1.6562500000000002, 1974574552],
+    [3.125, 1.5937500000000004, 1974574552],
+    [2.71875, 1.4687499999999998, 1974574552],
+    [2.40625, 1.3437500000000002, 1974574552],
+    [2.0625, 1.2187499999999996, 1974574552],
+    [2.5, 1.0312500000000002, 1974574552],
+    [2.84375, 1.0312500000000002, 1974574552],
+    [3.375, 0.9999999999999998, 1974574552],
+    [3.96875, 0.9687500000000004, 1974574552],
+    [4.6875, 0.9375, 1974574552],
+    [5.21875, 0.9375, 1974574552],
+    [5.84375, 0.9687500000000004, 1974574552],
+    [6.4375, 1.0312500000000002, 1974574552],
+    [7, 1.1562499999999998, 1974574552],
+    [7.34375, 1.2187499999999996, 1974574552],
+    [7.40625, 1.25, 1974574552],
+]
+
+# Robot Path:
+robot_path = [
+    [2.5, 0.2500000000000002, 1097479840],
+    [3.03125, 0.2500000000000002, 1097479840],
+    [3.625, 0.2500000000000002, 1097479840],
+    [4.125, 0.2500000000000002, 1097479840],
+    [4.625, 0.2500000000000002, 1097479840],
+    [5.375, 0.21874999999999978, 1097479840],
+    [5.90625, 0.18750000000000044, 1097479840],
+    [6.4375, 0.18750000000000044, 1097479840],
+    [6.96875, 0.18750000000000044, 1097479840],
+    [7.46875, 0.28124999999999956, 1097479840],
+    [7.9375, 0.46875, 1097479840],
+    [8, 0.9999999999999998, 1097479840],
+    [7.65625, 1.40625, 1097479840],
+    [7.15625, 1.5625, 1097479840],
+    [6.5625, 1.5625, 1097479840],
+    [6, 1.5312499999999996, 1097479840],
+    [5.5, 1.5000000000000002, 1097479840],
+    [4.9375, 1.40625, 1097479840],
+    [4.40625, 1.3749999999999996, 1097479840],
+    [4.9375, 1.3437500000000002, 1097479840],
+    [5.46875, 1.4375000000000004, 1097479840],
+    [6.0625, 1.4375000000000004, 1097479840],
+    [6.625, 1.4687499999999998, 1097479840],
+    [7.1875, 1.5000000000000002, 1097479840],
+    [7.65625, 1.71875, 1097479840],
+    [7.8125, 2.2187500000000004, 1097479840],
+    [7.71875, 2.71875, 1097479840],
+    [7.25, 2.90625, 1097479840],
+    [6.6875, 3.03125, 1097479840],
+    [6.03125, 3.125, 1097479840],
+    [5.4375, 3.125, 1097479840],
+    [4.875, 3.0625, 1097479840],
+    [4.28125, 3.03125, 1097479840],
+    [3.71875, 3.0000000000000004, 1097479840],
+    [4.34375, 2.96875, 1097479840],
+    [4.96875, 2.96875, 1097479840],
+    [5.65625, 3.0000000000000004, 1097479840],
+    [6.25, 3.0000000000000004, 1097479840],
+    [7.09375, 3.0000000000000004, 1097479840],
+    [7.75, 3.03125, 1097479840],
+    [8.25, 3.125, 1097479840],
+    [8.1875, 3.6875, 1097479840],
+    [7.59375, 4, 1097479840],
+    [6.96875, 4.21875, 1097479840],
+    [6.375, 4.375, 1097479840],
+    [5.875, 4.4375, 1097479840],
+    [5.375, 4.5, 1097479840],
+    [6, 4.5625, 1097479840],
+    [6.53125, 4.6875, 1097479840],
+    [7.0625, 4.78125, 1097479840],
+    [7.4375, 5.1875, 1097479840],
+    [7.1875, 5.6875, 1097479840],
+    [6.75, 6.03125, 1097479840],
+    [6.125, 6.34375, 1097479840],
+    [5.59375, 6.46875, 1097479840],
+    [5.0625, 6.46875, 1097479840],
+    [5.59375, 6.46875, 1097479840],
+    [6.78125, 6.5625, 1097479840],
+    [7.3125, 6.65625, 1097479840],
+    [7.9375, 6.90625, 1097479840],
+    [8.15625, 7.40625, 1097479840],
+    [7.6875, 7.71875, 1097479840],
+    [7.09375, 7.90625, 1097479840],
+    [6.46875, 8.0625, 1097479840],
+    [5.96875, 8.15625, 1097479840],
+    [5.375, 8.1875, 1097479840],
+    [4.6875, 8.3125, 1097479840],
+    [4.1875, 8.34375, 1097479840],
+    [4.6875, 8.4375, 1097479840],
+    [5.46875, 8.46875, 1097479840],
+    [5.96875, 8.46875, 1097479840],
+    [6.53125, 8.46875, 1097479840],
+    [7.09375, 8.46875, 1097479840],
+    [7.59375, 8.53125, 1097479840],
+    [8.15625, 8.625, 1097479840],
+    [8.65625, 8.8125, 1097479840],
+    [8.28125, 9.34375, 1097479840],
+    [7.75, 9.4375, 1097479840],
+    [7.15625, 9.5, 1097479840],
+    [6.625, 9.5, 1097479840],
+    [6, 9.5, 1097479840],
+    [5.375, 9.53125, 1097479840],
+    [4.8125, 9.53125, 1097479840],
+    [4.1875, 9.53125, 1097479840],
+    [3.5625, 9.4375, 1097479840],
+    [3, 9.4375, 1097479840],
+    [2.46875, 9.40625, 1097479840],
+    [1.84375, 9.34375, 1097479840],
+    [1.4375, 9.34375, 1097479840],
+]
+
+
+def perturb_walls(w, idx, amount=0.1):
+    # Add small random offset to one wall endpoint
+    w[idx][1] = w[idx][1] + amount
+    return w
+
+
+def get_robot(p_noise):
+    return RobotCapabilities(
+        p_noise=jnp.array(p_noise),
+        hd_noise=jnp.array(0.03),
+        sensor_noise=jnp.array(0.1),
+        n_sensors=jnp.array(8),
+        sensor_range=jnp.array(10.0),
+    )
+
+
+# Test 2: Perturb robot capabilities
+
+
+# Create random keys
+keys = jax.random.split(PRNGKey(0), 100)
+
+# Test 1: Run simulation with perturbed walls
+print("Test 1: Perturbing walls")
+start = time.time()
+all_paths_1, all_readings_1 = jax.vmap(
+    lambda k: simulate_robot(
+        World(*walls_to_jax(walls)),
+        get_robot(0.1),
+        jnp.array(robot_path),
+        k,
+    )
+)(keys)
+print(f"Wall perturbation test took {time.time() - start:.3f} seconds")
+
+# Test 2: Run simulation with perturbed robot capabilities
+print("\nTest 2: Perturbing robot capabilities")
+start = time.time()
+p_noises = jnp.linspace(0.05, 0.15, 100)  # Range of p_noise values
+all_paths_2, all_readings_2 = jax.vmap(
+    lambda k, p: simulate_robot(
+        World(*walls_to_jax(walls)), get_robot(p), jnp.array(robot_path), k
+    )
+)(keys, p_noises)
+print(f"Robot capability perturbation test took {time.time() - start:.3f} seconds")