Add support for jaxsim simulator with examples #20

xela-95 · 2024-10-10T14:19:31Z

In this PR I'm adding support for JaxSim simulator and explaining how to use its interface with a notebook similar to the ones already available for the other simulators.

flferretti

Early comments.

They can be rapidly addressed using pyupgrade as pyupgrade src/comodo/jaxsimSimulator/jaxsimSimulator.py

src/comodo/jaxsimSimulator/jaxsimSimulator.py

diegoferigo

Early review.

src/comodo/jaxsimSimulator/jaxsimSimulator.py

xela-95 · 2024-10-11T16:23:59Z

I'm waiting for #21 to be fixed and be able to run on a server the same notebook, but the implementation is now complete.

examples/jaxsim_walking.ipynb

diegoferigo

Thanks @xela-95 for this PR! I left minor nits (non-blocking). I'm facing the following problem when running the notebook, does it reproduce also on other setups?

AttributeError                            Traceback (most recent call last)
Cell In[6], line 15
     13 step_length = 0.1
     14 mpc = CentroidalMPC(robot_model=robot_model_init, step_length=step_length)
---> 15 mpc.intialize_mpc(mpc_parameters=mpc_parameters)
     17 # Set desired quantities
     18 mpc.configure(s_init=s_0, H_b_init=H_b_0)

File ~/git/comodo/src/comodo/centroidalMPC/centroidalMPC.py:93, in CentroidalMPC.intialize_mpc(self, mpc_parameters)
     91 self.contact_0_handler.set_parameter_int("number_of_corners", 4)
     92 self.contact_0_handler.set_parameter_string("contact_name", "left_foot")
---> 93 self.contact_0_handler.set_parameter_vector_float("corner_0",self.robot_model.corner_0)
     94 self.contact_0_handler.set_parameter_vector_float("corner_1", self.robot_model.corner_1)
     95 self.contact_0_handler.set_parameter_vector_float(
     96     "corner_2", self.robot_model.corner_2
     97 )

AttributeError: 'RobotModel' object has no attribute 'corner_0'

src/comodo/jaxsimSimulator/jaxsimSimulator.py

flferretti · 2024-10-14T08:17:22Z

Thanks @xela-95 for this PR! I left minor nits (non-blocking). I'm facing the following problem when running the notebook, does it reproduce also on other setups?

AttributeError                            Traceback (most recent call last)
Cell In[6], line 15
     13 step_length = 0.1
     14 mpc = CentroidalMPC(robot_model=robot_model_init, step_length=step_length)
---> 15 mpc.intialize_mpc(mpc_parameters=mpc_parameters)
     17 # Set desired quantities
     18 mpc.configure(s_init=s_0, H_b_init=H_b_0)

File ~/git/comodo/src/comodo/centroidalMPC/centroidalMPC.py:93, in CentroidalMPC.intialize_mpc(self, mpc_parameters)
     91 self.contact_0_handler.set_parameter_int("number_of_corners", 4)
     92 self.contact_0_handler.set_parameter_string("contact_name", "left_foot")
---> 93 self.contact_0_handler.set_parameter_vector_float("corner_0",self.robot_model.corner_0)
     94 self.contact_0_handler.set_parameter_vector_float("corner_1", self.robot_model.corner_1)
     95 self.contact_0_handler.set_parameter_vector_float(
     96     "corner_2", self.robot_model.corner_2
     97 )

AttributeError: 'RobotModel' object has no attribute 'corner_0'

Should be fixed by #20 (comment)

xela-95 · 2024-10-14T08:18:02Z

Thanks @xela-95 for this PR! I left minor nits (non-blocking). I'm facing the following problem when running the notebook, does it reproduce also on other setups?

Thanks for the review @diegoferigo! Since I detached this branch from a previous fork of @flferretti, I built my code on his changes, and some of them were needed for this notebook to run correctly (I had these local changes active on my machine). I will update the PR soon to include the missing changes to comodo.

diegoferigo · 2024-10-14T08:33:01Z

After the fix reported in #20 (comment), FYI I still get the following error, likely related to the changes you still need to port here:

In [22]: 
    ...: 
    ...: now = time.perf_counter()
    ...: 
    ...: logs = simulate(T=T, js=js, tsid=tsid, mpc=mpc, s_ref=s_0)
    ...: 
    ...: wall_time = time.perf_counter() - now
    ...: avg_iter_time_ms = (wall_time / (T / js_dt)) * 1000
    ...: 
    ...: print(
    ...:     f"\nRunning simulation took {wall_time:.2f}s for {T:.3f}s simulated time. \nIteration avg time of {avg_iter_time_ms:.1f} ms."
    ...: )
    ...: print(f"RTF: {T / wall_time * 100:.2f}%")
n_step_mpc_tsid=10, n_step_tsid_js=20
Exception during simulation at time0.0: 'CentroidalMPC' object has no attribute 'com_traj'
Traceback (most recent call last):
  File "<ipython-input-21-57d7e1812ec4>", line 52, in simulate
    mpc.update_references()
  File "/home/dferigo/git/comodo/src/comodo/centroidalMPC/centroidalMPC.py", line 215, in update_references
    self.com_traj = self.com_traj[1:]
AttributeError: 'CentroidalMPC' object has no attribute 'com_traj'

Running simulation took 0.00s for 6.000s simulated time. 
Iteration avg time of 0.0 ms.
RTF: 347104.11%

README.md

examples/jaxsim_walking.ipynb

diegoferigo

I successfully ran the notebook locally with the RigidContacts model 🎉 I still have to try with other contact models, but for the moment this PR is good to go from my side. We can enhance it as needed in the future. Thanks @xela-95!

src/comodo/jaxsimSimulator/jaxsimSimulator.py

xela-95 · 2024-10-22T08:00:13Z

@CarlottaSartore this PR is ready for your review :)

CarlottaSartore

Looks good to me, If it is possible I would divide the loading of the model from the initialization of the simulator

src/comodo/jaxsimSimulator/jaxsimSimulator.py

README.md

Co-authored-by: Filippo Ferretti <[email protected]>

Simplify conda environment installation Co-authored-by: Filippo Luca Ferretti <[email protected]>

Used `pyupgrade` for this. Co-authored-by: Filippo Ferretti <[email protected]>

Co-authored-by: Diego Ferigo <[email protected]> Co-authored-by: Filippo Luca Ferretti <[email protected]>

- update enum class for contact model type - manage the different ordering of joints internally (as done in mujoco) - Add docstrings for public apis - import contact models from the rbda.contacts module - Added a new property to track simulation time - Modified the step() method to support dry run mode for simulation - Updated the rendering logic to use the new simulation time property - Optimized the update of link contact forces - Added a new property to retrieve the contact model type

Addressing PR review comments. Co-authored-by: Diego Ferigo <[email protected]> Co-authored-by: Filippo Luca Ferretti <[email protected]>

Co-authored-by: Filippo Ferretti <[email protected]>

…ad_model` to the constructor

CarlottaSartore · 2024-10-23T13:10:10Z

Great @xela-95 thank you ! Merging !

xela-95 added enhancement New feature or request phase:implementation labels Oct 10, 2024

xela-95 requested review from diegoferigo, CarlottaSartore and flferretti October 10, 2024 14:19

xela-95 self-assigned this Oct 10, 2024

xela-95 marked this pull request as draft October 10, 2024 14:27

flferretti reviewed Oct 10, 2024

View reviewed changes

diegoferigo suggested changes Oct 10, 2024

View reviewed changes

xela-95 mentioned this pull request Oct 11, 2024

OSError when creating RobotModel #21

Closed

xela-95 marked this pull request as ready for review October 11, 2024 16:24

flferretti reviewed Oct 11, 2024

View reviewed changes

examples/jaxsim_walking.ipynb Show resolved Hide resolved

diegoferigo suggested changes Oct 14, 2024

View reviewed changes

xela-95 force-pushed the add-jaxsim-simulator branch 2 times, most recently from 3ac21c8 to fcea6d9 Compare October 14, 2024 12:08

flferretti reviewed Oct 14, 2024

View reviewed changes

README.md Outdated Show resolved Hide resolved

flferretti reviewed Oct 14, 2024

View reviewed changes

examples/jaxsim_walking.ipynb Outdated Show resolved Hide resolved

diegoferigo approved these changes Oct 14, 2024

View reviewed changes

xela-95 mentioned this pull request Oct 15, 2024

Add pixi support #23

Merged

diegoferigo reviewed Oct 21, 2024

View reviewed changes

src/comodo/jaxsimSimulator/jaxsimSimulator.py Show resolved Hide resolved

flferretti reviewed Oct 21, 2024

View reviewed changes

src/comodo/jaxsimSimulator/jaxsimSimulator.py Outdated Show resolved Hide resolved

CarlottaSartore approved these changes Oct 22, 2024

View reviewed changes

src/comodo/jaxsimSimulator/jaxsimSimulator.py Show resolved Hide resolved

README.md Show resolved Hide resolved

xela-95 and others added 3 commits October 23, 2024 10:09

Add JaxsimSimulator class

0b87a45

Co-authored-by: Filippo Ferretti <[email protected]>

Add __init__.py

d840503

Update gitignore

2ee50af

xela-95 and others added 15 commits October 23, 2024 10:09

Add jaxsim_walking example notebook

d0df93d

Update README

b753d99

Simplify conda environment installation Co-authored-by: Filippo Luca Ferretti <[email protected]>

Upgrade to modern python conventions

5faf62d

Used `pyupgrade` for this. Co-authored-by: Filippo Ferretti <[email protected]>

Apply suggestions from code review

a15c8e0

Co-authored-by: Diego Ferigo <[email protected]> Co-authored-by: Filippo Luca Ferretti <[email protected]>

Update __init__.py

9778b4f

Update verbosity level in TSIDController.py

cfbe875

Update jaxsim_walking.ipynb

2711f8f

Refactor link_contact_forces property in jaxsimSimulator.py

4bb884f

Update CentroidalMPC parameters

a8dc7f4

Update jaxsimSimulator.py

29c70b7

Addressing PR review comments. Co-authored-by: Diego Ferigo <[email protected]> Co-authored-by: Filippo Luca Ferretti <[email protected]>

Update jaxsim_walking.ipynb

256b828

Co-authored-by: Filippo Ferretti <[email protected]>

Update contact model parameters in JaxsimSimulator.py

74cf682

Update jaxsimSimulator by moving dt and contact_model_type from `lo…

ea094bd

…ad_model` to the constructor

Update JaxsimSimulator initialization in jaxsim_walking.ipynb

28c219d

xela-95 force-pushed the add-jaxsim-simulator branch from 074acf6 to 28c219d Compare October 23, 2024 08:14

CarlottaSartore merged commit 51e43e4 into ami-iit:main Oct 23, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for jaxsim simulator with examples #20

Add support for jaxsim simulator with examples #20

xela-95 commented Oct 10, 2024

flferretti left a comment •

edited

Loading

diegoferigo left a comment

xela-95 commented Oct 11, 2024

diegoferigo left a comment

flferretti commented Oct 14, 2024

xela-95 commented Oct 14, 2024

diegoferigo commented Oct 14, 2024

diegoferigo left a comment

xela-95 commented Oct 22, 2024

CarlottaSartore left a comment

CarlottaSartore commented Oct 23, 2024

Add support for jaxsim simulator with examples #20

Add support for jaxsim simulator with examples #20

Conversation

xela-95 commented Oct 10, 2024

flferretti left a comment • edited Loading

Choose a reason for hiding this comment

diegoferigo left a comment

Choose a reason for hiding this comment

xela-95 commented Oct 11, 2024

diegoferigo left a comment

Choose a reason for hiding this comment

flferretti commented Oct 14, 2024

xela-95 commented Oct 14, 2024

diegoferigo commented Oct 14, 2024

diegoferigo left a comment

Choose a reason for hiding this comment

xela-95 commented Oct 22, 2024

CarlottaSartore left a comment

Choose a reason for hiding this comment

CarlottaSartore commented Oct 23, 2024

flferretti left a comment •

edited

Loading