-
Notifications
You must be signed in to change notification settings - Fork 9
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.
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
Introduce new functional APIs #88
Conversation
b12bc31
to
1903889
Compare
The MWE of the new functional APIs is the following script. Resulting video attached. MWEimport pathlib
import jax.numpy as jnp
import jaxsim.api as js
import numpy as np
import resolve_robotics_uri_py
import rod
from jaxsim import VelRepr, integrators
# Find the urdf file.
urdf_path = resolve_robotics_uri_py.resolve_robotics_uri(
uri="model://ergoCubSN001/model.urdf"
)
# Build the ROD model.
rod_sdf = rod.Sdf.load(sdf=urdf_path)
# Build the model.
model = js.model.JaxSimModel.build_from_model_description(
model_description=rod_sdf.model,
gravity=jnp.array([0, 0, -10.0]),
)
# Reduce the model.
model = js.model.reduce(
model=model,
considered_joints=tuple(
[
j
for j in model.joint_names()
if "camera" not in j
and "neck" not in j
and "wrist" not in j
and "thumb" not in j
and "index" not in j
and "middle" not in j
and "ring" not in j
and "pinkie" not in j
]
),
)
# Build the model's data.
# Set already here the initial base position.
data0 = js.data.JaxSimModelData.build(
model=model,
base_position=jnp.array([0, 0, 0.85]),
velocity_representation=VelRepr.Inertial,
)
# Update the soft-contact parameters.
# By default, only 1 support point is used as worst-case scenario.
# Feel free to tune this with more points to get a less stiff system.
data0 = data0.replace(
soft_contacts_params=js.contact.estimate_good_soft_contacts_parameters(
model, number_of_active_collidable_points_steady_state=2
)
)
# =====================
# Create the integrator
# =====================
# Create a RK4 integrator integrating the quaternion on SO(3).
integrator = integrators.fixed_step.RungeKutta4SO3.build(
dynamics=js.ode.wrap_system_dynamics_for_integration(
model=model,
data=data0,
system_dynamics=js.ode.system_dynamics,
),
)
# =========================================
# Visualization in Mujoco viewer / renderer
# =========================================
from jaxsim.mujoco import MujocoModelHelper, MujocoVideoRecorder, RodModelToMjcf
# Convert the ROD model to a Mujoco model.
mjcf_string, assets = RodModelToMjcf.convert(
rod_model=rod_sdf.models()[0],
considered_joints=list(model.joint_names()),
)
# Build the Mujoco model helper.
mj_model_helper = self = MujocoModelHelper.build_from_xml(
mjcf_description=mjcf_string, assets=assets
)
# Create the video recorder.
recorder = MujocoVideoRecorder(
model=mj_model_helper.model,
data=mj_model_helper.data,
fps=int(1 / 0.010),
width=320 * 4,
height=240 * 4,
)
# ==============
# Recording loop
# ==============
# Initialize the integrator.
t0 = 0.0
tf = 5.0
dt = 0.001_000
integrator_state = integrator.init(x0=data0.state, t0=t0, dt=dt)
# Initialize the loop.
data = data0.copy()
joint_names = list(model.joint_names())
while data.time_ns < tf * 1e9:
# Integrate the dynamics.
data, integrator_state = js.model.step(
dt=dt,
model=model,
data=data,
integrator=integrator,
integrator_state=integrator_state,
# Optional inputs
joint_forces=None,
external_forces=None,
)
# Extract the generalized position.
s = data.state.physics_model.joint_positions
W_p_B = data.state.physics_model.base_position
W_Q_B = data.state.physics_model.base_quaternion
# Update the data object stored in the helper, which is shared with the recorder.
mj_model_helper.set_base_position(position=np.array(W_p_B))
mj_model_helper.set_base_orientation(orientation=np.array(W_Q_B), dcm=False)
mj_model_helper.set_joint_positions(positions=np.array(s), joint_names=joint_names)
# Record the frame if the time is right to get the desired fps.
if data.time_ns % jnp.array(1e9 / recorder.fps).astype(jnp.uint64) == 0:
recorder.render_frame(camera_name=None)
# Store the video.
video_path = pathlib.Path("~/video.mp4").expanduser()
recorder.write_video(path=video_path, exist_ok=True)
# Clean up the recorder.
recorder.frames = []
recorder.renderer.close() VideoResulting soft-contacts parameters: SoftContactsParams(
K=Array(465529.79648, dtype=float64),
D=Array(10270.41887, dtype=float64),
mu=Array(0.5, dtype=float64)
) video.mp4Some comment:
|
bd8b164
to
5965f1c
Compare
Thanks for working on this @diegoferigo, it looks pretty neat! Could you please add the new modules in the docs? |
8477d1b
to
ac1532c
Compare
I don't want to waste time now to fix the documentation, I'll do it as soon as things get more stable. I expect problem with math rendering etc, and I don't want to dedicate time to it at this early stage :) It has to be done before the next release. |
Agreed! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks a lot Diego, LGTM
I have a couple more comments:
data.state.physics_model.replace(
joint_positions=new_position
) C.C. @traversaro |
I already have local tests, but not yet ready. They will be added in a new PR.
Although the new functional APIs will get merged into |
ac1532c
to
5ad41dd
Compare
@flferretti I just added |
3798121
to
74eea83
Compare
Co-authored-by: Filippo Luca Ferretti <[email protected]>
Co-authored-by: Filippo Luca Ferretti <[email protected]>
74eea83
to
e5dc210
Compare
Great! |
These new functional APIs (
jaxsim.api
) will soon replace those based on OOP (jaxsim.high_level
).jax.vmap
andjax.grad
is more simple being more vanilla.Furthermore, this PR:
jax.vmap
to iterate on links and joints.Some implementation detail that might ease the review process and the transition to the new APIs:
JaxSimModel
andJaxSimModelData
.kwargs
.Next steps:
jaxsim.api.model.JaxSimModel
withjaxsim.physics.model.PhysicsModel
.KynDynParameters
to store insideJaxSimModel
the model parameters, so that we can differentiate against them as well (we need also decent setters/getters).📚 Documentation preview 📚: https://jaxsim--88.org.readthedocs.build//88/