From 8dbee8d566722a4acc809ddc648d7553e9cbc8b8 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sat, 24 Feb 2024 13:24:59 -0500 Subject: [PATCH 1/5] minor changes --- python/mujoco_mpc/mjx/tasks/bimanual/handover.py | 5 ++--- python/mujoco_mpc/mjx/visualize.py | 6 +++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py index ccbcfd24e..da6d1c6c4 100644 --- a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py +++ b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py @@ -18,8 +18,7 @@ from jax import numpy as jp import mujoco from mujoco import mjx -from mujoco_mpc.mjx import predictive_sampling - +from typing import Callable def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: """Returns cost for bimanual bring to target task.""" @@ -48,7 +47,7 @@ def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: def get_models_and_cost_fn() -> ( - tuple[mujoco.MjModel, mujoco.MjModel, predictive_sampling.CostFn] + tuple[mujoco.MjModel, mujoco.MjModel, Callable[[mjx.Model, mjx.Data], jax.Array]] ): """Returns a tuple of the model and the cost function.""" path = epath.Path( diff --git a/python/mujoco_mpc/mjx/visualize.py b/python/mujoco_mpc/mjx/visualize.py index deb493d65..439657eae 100644 --- a/python/mujoco_mpc/mjx/visualize.py +++ b/python/mujoco_mpc/mjx/visualize.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - +# %% import matplotlib.pyplot as plt import mediapy import mujoco -from mujoco_mpc.mjx import predictive_sampling -from mujoco_mpc.mjx.tasks.bimanual import handover +import predictive_sampling +from tasks.bimanual import handover # %% nsteps = 500 steps_per_plan = 4 From 8e55b3b5cdf1d5db9e0144fa6077774f580eb2d5 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sat, 24 Feb 2024 15:29:01 -0500 Subject: [PATCH 2/5] setup aloha mjx cmake --- CMakeLists.txt | 2 +- mjpc/tasks/CMakeLists.txt | 24 +++++++++++++++++++ .../mujoco_mpc/mjx/tasks/bimanual/handover.py | 23 +++++++----------- 3 files changed, 33 insertions(+), 16 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 51c276599..23d7dff04 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -61,7 +61,7 @@ set(MUJOCO_MPC_MUJOCO_GIT_TAG ) set(MUJOCO_MPC_MENAGERIE_GIT_TAG - 8ef01e87fffaa8ec634a4826c5b2092733b2f3c8 + 8a5f659ac3607dc5adb988e0187f683fe0f4edf4 CACHE STRING "Git revision for MuJoCo Menagerie." ) diff --git a/mjpc/tasks/CMakeLists.txt b/mjpc/tasks/CMakeLists.txt index 72a7cb358..3cbfa14a7 100644 --- a/mjpc/tasks/CMakeLists.txt +++ b/mjpc/tasks/CMakeLists.txt @@ -162,6 +162,30 @@ add_custom_target( ${CMAKE_CURRENT_BINARY_DIR}/bimanual/aloha.xml <${CMAKE_CURRENT_SOURCE_DIR}/bimanual/aloha.patch + # MJX ALOHA + COMMAND ${CMAKE_COMMAND} -E copy + ${menagerie_SOURCE_DIR}/aloha/mjx_aloha.patch + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_aloha.patch + COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_aloha.xml + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/aloha.xml + <${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_aloha.patch + COMMAND ${CMAKE_COMMAND} -E copy + ${menagerie_SOURCE_DIR}/aloha/scene.xml + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/scene.xml + COMMAND ${CMAKE_COMMAND} -E copy + ${menagerie_SOURCE_DIR}/aloha/mjx_scene.patch + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_scene.patch + COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_scene.xml + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/scene.xml + <${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_scene.patch + + COMMAND ${CMAKE_COMMAND} -E copy + ${menagerie_SOURCE_DIR}/aloha/mjx_integrated_cartesian_actuators.patch + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_integrated_cartesian_actuators.patch + COMMAND patch -o ${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_integrated_cartesian_actuators.xml + ${CMAKE_CURRENT_BINARY_DIR}/bimanual/integrated_cartesian_actuators.xml + <${CMAKE_CURRENT_BINARY_DIR}/bimanual/mjx_integrated_cartesian_actuators.patch + COMMAND ${CMAKE_COMMAND} -E copy ${menagerie_SOURCE_DIR}/robotis_op3/op3.xml ${CMAKE_CURRENT_BINARY_DIR}/op3/op3.xml diff --git a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py index da6d1c6c4..9b0774711 100644 --- a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py +++ b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py @@ -13,11 +13,12 @@ # limitations under the License. # ============================================================================== -from etils import epath import jax from jax import numpy as jp import mujoco from mujoco import mjx +import os +import pathlib from typing import Callable def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: @@ -50,19 +51,11 @@ def get_models_and_cost_fn() -> ( tuple[mujoco.MjModel, mujoco.MjModel, Callable[[mjx.Model, mjx.Data], jax.Array]] ): """Returns a tuple of the model and the cost function.""" - path = epath.Path( - 'build/mujoco_menagerie/aloha/' - ) - model_file_name = 'mjx_scene.xml' - xml = (path / model_file_name).read_text() - assets = {} - for f in path.glob('*.xml'): - if f.name == model_file_name: - continue - assets[f.name] = f.read_bytes() - for f in (path / 'assets').glob('*'): - assets[f.name] = f.read_bytes() - sim_model = mujoco.MjModel.from_xml_string(xml, assets) - plan_model = mujoco.MjModel.from_xml_string(xml, assets) + path = ( + pathlib.Path(os.path.abspath("")).parent.parent.parent + / "build/mjpc/tasks/bimanual/mjx_scene.xml" + ) + sim_model = mujoco.MjModel.from_xml_path(str(path)) + plan_model = mujoco.MjModel.from_xml_path(str(path)) plan_model.opt.timestep = 0.01 # incidentally, already the case return sim_model, plan_model, bring_to_target From da7be0c044d331058ac1fd624499722ed79a4004 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sat, 24 Feb 2024 15:46:50 -0500 Subject: [PATCH 3/5] minor fixes --- python/mujoco_mpc/mjx/tasks/bimanual/handover.py | 4 +++- python/mujoco_mpc/mjx/visualize.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py index 9b0774711..81a163788 100644 --- a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py +++ b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py @@ -21,6 +21,8 @@ import pathlib from typing import Callable +CostFn = Callable[[mjx.Model, mjx.Data], jax.Array] + def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: """Returns cost for bimanual bring to target task.""" # reach @@ -48,7 +50,7 @@ def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: def get_models_and_cost_fn() -> ( - tuple[mujoco.MjModel, mujoco.MjModel, Callable[[mjx.Model, mjx.Data], jax.Array]] + tuple[mujoco.MjModel, mujoco.MjModel, CostFn] ): """Returns a tuple of the model and the cost function.""" path = ( diff --git a/python/mujoco_mpc/mjx/visualize.py b/python/mujoco_mpc/mjx/visualize.py index 439657eae..61ab891df 100644 --- a/python/mujoco_mpc/mjx/visualize.py +++ b/python/mujoco_mpc/mjx/visualize.py @@ -19,7 +19,7 @@ import predictive_sampling from tasks.bimanual import handover # %% -nsteps = 500 +nsteps = 1 steps_per_plan = 4 frame_skip = 5 # how many steps between each rendered frame @@ -34,7 +34,7 @@ nsample=128 - 1, interp='zero', ) - +# %% trajectory, costs, plan_time = ( predictive_sampling.receding_horizon_optimization( p, From 34342874432e4b4160a6aae3d0e20486867b5917 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sun, 25 Feb 2024 19:32:22 -0500 Subject: [PATCH 4/5] restore mjx example --- .../mujoco_mpc/mjx/tasks/bimanual/handover.py | 30 +++++++++++-------- python/mujoco_mpc/mjx/visualize.py | 12 ++++---- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py index 81a163788..780535f92 100644 --- a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py +++ b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py @@ -13,15 +13,13 @@ # limitations under the License. # ============================================================================== +from etils import epath import jax from jax import numpy as jp import mujoco from mujoco import mjx -import os -import pathlib -from typing import Callable +from mujoco_mpc.mjx import predictive_sampling -CostFn = Callable[[mjx.Model, mjx.Data], jax.Array] def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: """Returns cost for bimanual bring to target task.""" @@ -50,14 +48,22 @@ def bring_to_target(m: mjx.Model, d: mjx.Data) -> jax.Array: def get_models_and_cost_fn() -> ( - tuple[mujoco.MjModel, mujoco.MjModel, CostFn] + tuple[mujoco.MjModel, mujoco.MjModel, predictive_sampling.CostFn] ): """Returns a tuple of the model and the cost function.""" - path = ( - pathlib.Path(os.path.abspath("")).parent.parent.parent - / "build/mjpc/tasks/bimanual/mjx_scene.xml" - ) - sim_model = mujoco.MjModel.from_xml_path(str(path)) - plan_model = mujoco.MjModel.from_xml_path(str(path)) + path = epath.Path( + 'build/mjpc/tasks/bimanual/' + ) + model_file_name = 'mjx_scene.xml' + xml = (path / model_file_name).read_text() + assets = {} + for f in path.glob('*.xml'): + if f.name == model_file_name: + continue + assets[f.name] = f.read_bytes() + for f in (path / 'assets').glob('*'): + assets[f.name] = f.read_bytes() + sim_model = mujoco.MjModel.from_xml_string(xml, assets) + plan_model = mujoco.MjModel.from_xml_string(xml, assets) plan_model.opt.timestep = 0.01 # incidentally, already the case - return sim_model, plan_model, bring_to_target + return sim_model, plan_model, bring_to_target \ No newline at end of file diff --git a/python/mujoco_mpc/mjx/visualize.py b/python/mujoco_mpc/mjx/visualize.py index 61ab891df..014c5b061 100644 --- a/python/mujoco_mpc/mjx/visualize.py +++ b/python/mujoco_mpc/mjx/visualize.py @@ -12,14 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -# %% + import matplotlib.pyplot as plt import mediapy import mujoco -import predictive_sampling -from tasks.bimanual import handover +from mujoco_mpc.mjx import predictive_sampling +from mujoco_mpc.mjx.tasks.bimanual import handover # %% -nsteps = 1 +nsteps = 500 steps_per_plan = 4 frame_skip = 5 # how many steps between each rendered frame @@ -34,7 +34,7 @@ nsample=128 - 1, interp='zero', ) -# %% + trajectory, costs, plan_time = ( predictive_sampling.receding_horizon_optimization( p, @@ -70,4 +70,4 @@ frames.append(renderer.render()) # %% mediapy.show_video(frames, fps=1/sim_model.opt.timestep/frame_skip) -# %% +# %% \ No newline at end of file From 59b86b84d058fc5fb5108d72dbd487a5ce803b4e Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sun, 25 Feb 2024 19:33:05 -0500 Subject: [PATCH 5/5] endline --- python/mujoco_mpc/mjx/tasks/bimanual/handover.py | 2 +- python/mujoco_mpc/mjx/visualize.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py index 780535f92..4d1d81205 100644 --- a/python/mujoco_mpc/mjx/tasks/bimanual/handover.py +++ b/python/mujoco_mpc/mjx/tasks/bimanual/handover.py @@ -66,4 +66,4 @@ def get_models_and_cost_fn() -> ( sim_model = mujoco.MjModel.from_xml_string(xml, assets) plan_model = mujoco.MjModel.from_xml_string(xml, assets) plan_model.opt.timestep = 0.01 # incidentally, already the case - return sim_model, plan_model, bring_to_target \ No newline at end of file + return sim_model, plan_model, bring_to_target diff --git a/python/mujoco_mpc/mjx/visualize.py b/python/mujoco_mpc/mjx/visualize.py index 014c5b061..deb493d65 100644 --- a/python/mujoco_mpc/mjx/visualize.py +++ b/python/mujoco_mpc/mjx/visualize.py @@ -70,4 +70,4 @@ frames.append(renderer.render()) # %% mediapy.show_video(frames, fps=1/sim_model.opt.timestep/frame_skip) -# %% \ No newline at end of file +# %%