diff --git a/scripts/experiments/icra/mug/mug.pynb b/scripts/experiments/icra/mug/mug.ipynb similarity index 82% rename from scripts/experiments/icra/mug/mug.pynb rename to scripts/experiments/icra/mug/mug.ipynb index a585d41a..afc8a2cc 100644 --- a/scripts/experiments/icra/mug/mug.pynb +++ b/scripts/experiments/icra/mug/mug.ipynb @@ -29,7 +29,7 @@ "output_type": "stream", "text": [ "You can open the visualizer by visiting the following URL:\n", - "http://127.0.0.1:7014/static/\n" + "http://127.0.0.1:7016/static/\n" ] } ], @@ -40,58 +40,6 @@ { "cell_type": "code", "execution_count": 3, - "id": "d705f5f3-499d-4f43-b689-41ba648a2a1f", - "metadata": {}, - "outputs": [], - "source": [ - "model_dir = os.path.join(b.utils.get_assets_dir(),\"ycb_video_models/models/025_mug\")\n", - "mesh = b.utils.mesh.load_mesh(os.path.join(model_dir, \"textured_simple.obj\"))\n" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "7405b61a-4464-4ac3-ae86-f9eecc0965eb", - "metadata": {}, - "outputs": [], - "source": [ - "b.clear()\n", - "num_colors = 2\n", - "colors = b.viz.distinct_colors(num_colors)\n", - "offset = b.transform_from_pos(jnp.array([0.003, 0.0,0.0]))\n", - "for i in range(num_colors):\n", - " b.show_trimesh(f\"{i}\", mesh, color=colors[i])\n", - " b.set_pose(f\"{i}\", offset @ b.transform_from_axis_angle(jnp.array([0.0, 0.0, 1.0]), i))" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "c5eab559-3fd7-4964-a903-9f1a50ad3842", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(Array([0.116966, 0.093075, 0.081384], dtype=float32),\n", - " Array([[1.0000000e+00, 0.0000000e+00, 0.0000000e+00, 8.2690008e-03],\n", - " [0.0000000e+00, 1.0000000e+00, 0.0000000e+00, 6.1250106e-04],\n", - " [0.0000000e+00, 0.0000000e+00, 1.0000000e+00, 1.6890001e-03],\n", - " [0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 1.0000000e+00]], dtype=float32))" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "b.utils.aabb(mesh.vertices)" - ] - }, - { - "cell_type": "code", - "execution_count": 6, "id": "a42856ac-4db2-43cd-a4e7-6b93bc550f12", "metadata": {}, "outputs": [ @@ -140,7 +88,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 4, "id": "f4648f31-caf1-4792-bd83-9d652a8c5e4b", "metadata": {}, "outputs": [], @@ -156,7 +104,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 5, "id": "3277a542-3698-40f7-a998-d7febb9591eb", "metadata": {}, "outputs": [], @@ -251,145 +199,845 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 6, "id": "4a64e88e-4df4-4ae4-bd21-24c04fdc7782", "metadata": {}, "outputs": [ { - "data": { - "text/plain": [ - "['002_master_chef_can',\n", - " '003_cracker_box',\n", - " '004_sugar_box',\n", - " '005_tomato_soup_can',\n", - " '006_mustard_bottle',\n", - " '007_tuna_fish_can',\n", - " '008_pudding_box',\n", - " '009_gelatin_box',\n", - " '010_potted_meat_can',\n", - " '011_banana',\n", - " '019_pitcher_base',\n", - " '021_bleach_cleanser',\n", - " '024_bowl',\n", - " '025_mug',\n", - " '035_power_drill',\n", - " '036_wood_block',\n", - " '037_scissors',\n", - " '040_large_marker',\n", - " '051_large_clamp',\n", - " '052_extra_large_clamp',\n", - " '061_foam_brick']" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "b.utils.ycb_loader.MODEL_NAMES" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "772bb4b0-af42-4842-9af6-94752b30a707", - "metadata": {}, - "outputs": [], - "source": [ - "def get_depth_image(image):\n", - " mval = image[image < image.max()].max()\n", - " return b.get_depth_image(image, max=mval)\n", - "\n", - "def get_poses_non_jit(contact_params, id_table, id):\n", - " sg = b.scene_graph.SceneGraph(\n", - " root_poses=jnp.array([table_pose, jnp.eye(4)]),\n", - " box_dimensions=jnp.array([b.RENDERER.model_box_dims[id_table], b.RENDERER.model_box_dims[id]]),\n", - " parents=jnp.array([-1, 0]),\n", - " contact_params=jnp.array([jnp.zeros(3), contact_params]),\n", - " face_parent=jnp.array([-1,2]),\n", - " face_child=jnp.array([-1,3]),\n", - " )\n", - " poses = sg.get_poses()\n", - " return poses\n", - "get_poses = jax.jit(get_poses_non_jit)\n", - "\n", - "def render_image_non_jit(contact_params):\n", - " id_table = 21\n", - " id = 13\n", - " poses = get_poses_non_jit(contact_params, id_table, id)\n", - " img = b.RENDERER.render(\n", - " poses , jnp.array([id_table, id])\n", - " )[...,:3]\n", - " return img\n", - "render_image = jax.jit(render_image_non_jit)\n", - "\n", - "scorer = lambda obs, c, var, outlier_prob, outlier_volume: threedp3_likelihood(\n", - " obs, render_image_non_jit(c), var, outlier_prob, outlier_volume, 4\n", - ")\n", - "scorer_jit = jax.jit(scorer)\n", - "sweep_scorer = jax.jit(jax.vmap(\n", - " scorer \n", - " ,in_axes=(None, 0, None, None, None)))\n", - "\n", - "# sweep_scorer = jax.jit(jax.vmap(jax.vmap(jax.vmap(\n", - "# scorer, \n", - "# ,in_axes=(None, None, None, 0, None))\n", - "# ,in_axes=(None, None, 0, None, None))\n", - "# ,in_axes=(None, 0, None, None, None)))\n" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "fb05d62f-8b7f-4ed2-873c-530302c10964", - "metadata": {}, - "outputs": [], - "source": [ - "width = 0.01\n", - "ang = jnp.pi\n", - "contact_param_deltas = b.utils.make_translation_grid_enumeration_3d(\n", - " -width, -width, -ang,\n", - " width, width, ang,\n", - " 21,21,300\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "34bdcc5b-c781-4f35-bc8a-e89d35033465", - "metadata": {}, - "outputs": [], - "source": [ - "key = jax.random.PRNGKey(10)" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "id": "02d80355-07ff-49af-8233-71d276f95303", - "metadata": {}, - "outputs": [], - "source": [ - "variance = 0.0001\n", - "outlier_prob = 0.0001\n", - "outlier_volume = 1.0" - ] - }, - { - "cell_type": "code", - "execution_count": 18, - "id": "cae7a753-e4f7-4cc3-beb5-5d149e37119c", - "metadata": {}, - "outputs": [ + "data": { + "text/plain": [ + "['002_master_chef_can',\n", + " '003_cracker_box',\n", + " '004_sugar_box',\n", + " '005_tomato_soup_can',\n", + " '006_mustard_bottle',\n", + " '007_tuna_fish_can',\n", + " '008_pudding_box',\n", + " '009_gelatin_box',\n", + " '010_potted_meat_can',\n", + " '011_banana',\n", + " '019_pitcher_base',\n", + " '021_bleach_cleanser',\n", + " '024_bowl',\n", + " '025_mug',\n", + " '035_power_drill',\n", + " '036_wood_block',\n", + " '037_scissors',\n", + " '040_large_marker',\n", + " '051_large_clamp',\n", + " '052_extra_large_clamp',\n", + " '061_foam_brick']" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "b.utils.ycb_loader.MODEL_NAMES" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "772bb4b0-af42-4842-9af6-94752b30a707", + "metadata": {}, + "outputs": [], + "source": [ + "def get_depth_image(image):\n", + " mval = image[image < image.max()].max()\n", + " return b.get_depth_image(image, max=mval)\n", + "\n", + "def get_poses_non_jit(contact_params, id_table, id):\n", + " sg = b.scene_graph.SceneGraph(\n", + " root_poses=jnp.array([table_pose, jnp.eye(4)]),\n", + " box_dimensions=jnp.array([b.RENDERER.model_box_dims[id_table], b.RENDERER.model_box_dims[id]]),\n", + " parents=jnp.array([-1, 0]),\n", + " contact_params=jnp.array([jnp.zeros(3), contact_params]),\n", + " face_parent=jnp.array([-1,2]),\n", + " face_child=jnp.array([-1,3]),\n", + " )\n", + " poses = sg.get_poses()\n", + " return poses\n", + "get_poses = jax.jit(get_poses_non_jit)\n", + "\n", + "def render_image_non_jit(contact_params):\n", + " id_table = 21\n", + " id = 13\n", + " poses = get_poses_non_jit(contact_params, id_table, id)\n", + " img = b.RENDERER.render(\n", + " poses , jnp.array([id_table, id])\n", + " )[...,:3]\n", + " return img\n", + "render_image = jax.jit(render_image_non_jit)\n", + "\n", + "scorer = lambda obs, c, var, outlier_prob, outlier_volume: threedp3_likelihood(\n", + " obs, render_image_non_jit(c), var, outlier_prob, outlier_volume, 3\n", + ")\n", + "scorer_jit = jax.jit(scorer)\n", + "sweep_scorer = jax.jit(jax.vmap(\n", + " scorer \n", + " ,in_axes=(None, 0, None, None, None)))\n", + "\n", + "# sweep_scorer = jax.jit(jax.vmap(jax.vmap(jax.vmap(\n", + "# scorer, \n", + "# ,in_axes=(None, None, None, 0, None))\n", + "# ,in_axes=(None, None, 0, None, None))\n", + "# ,in_axes=(None, 0, None, None, None)))\n" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "fb05d62f-8b7f-4ed2-873c-530302c10964", + "metadata": {}, + "outputs": [], + "source": [ + "width = 0.015\n", + "ang = jnp.pi\n", + "contact_param_deltas = b.utils.make_translation_grid_enumeration_3d(\n", + " -width, -width, -ang,\n", + " width, width, ang,\n", + " 17,17,100\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "34bdcc5b-c781-4f35-bc8a-e89d35033465", + "metadata": {}, + "outputs": [], + "source": [ + "key = jax.random.PRNGKey(10)" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "02d80355-07ff-49af-8233-71d276f95303", + "metadata": {}, + "outputs": [], + "source": [ + "variance = 0.0001\n", + "outlier_prob = 0.001\n", + "outlier_volume = 1.0" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "cae7a753-e4f7-4cc3-beb5-5d149e37119c", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + " 0%| | 0/50 [00:00" ] @@ -974,53 +1622,14 @@ }, { "cell_type": "code", - "execution_count": 101, + "execution_count": null, "id": "e397b20d-889c-458a-be01-79cf19b509de", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "['002_master_chef_can',\n", - " '003_cracker_box',\n", - " '004_sugar_box',\n", - " '005_tomato_soup_can',\n", - " '006_mustard_bottle',\n", - " '007_tuna_fish_can',\n", - " '008_pudding_box',\n", - " '009_gelatin_box',\n", - " '010_potted_meat_can',\n", - " '011_banana',\n", - " '019_pitcher_base',\n", - " '021_bleach_cleanser',\n", - " '024_bowl',\n", - " '025_mug',\n", - " '035_power_drill',\n", - " '036_wood_block',\n", - " '037_scissors',\n", - " '040_large_marker',\n", - " '051_large_clamp',\n", - " '052_extra_large_clamp',\n", - " '061_foam_brick']" - ] - }, - "execution_count": 101, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "b.utils.ycb_loader.MODEL_NAMES" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "5e17f43d-833a-4984-ae79-9a1b262548f3", - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "code", "execution_count": null,