Using judo's viser GUI for development/trajopt purposes

[LaVieEstDure]

Hello! Big fan of judo, been considering making something similar myself. I've only had a brief look at the code but it seems that currently judo is built around mostly using it as an online and real time ROS-style node environment for deploying and testing sampling based MPC.

For both development/debugging sampling based MPC and for sampling based trajopt which doesn't require a real time sim, sometimes it's useful to have more raw access to the visualiser which is updated with whatever you feel like visualising in the algorithm loop, similar to mujoco's passive visualiser. However the mj passive visualiser is very low level, both graphically and visualisaiton data wise, and does not have the benefits of web based GUI (personally I can't get it to work over ssh due to Wayland issues with glfw).

While the Viser GUI class seems fairly well separated from the rest of the code to the point you could just import judo.visualizers.model and use it yourself with your own Viser loop, it might be a bit yucky to import internals and be prone to breakage in the future.

Is there plans to decouple the Viser GUI, make the API public facing and support a judo passive visualiser as an actual usecase?

[alberthli]

Hi @LaVieEstDure, your understanding is correct: our current support centers around the async/real-time case, because this closely resembles what real-world deployment of a sampling algorithm would look like.

That said, passive viewing with the viewing tools we have is actually really easy already! Here's an example:

import viser
from mujoco import MjData, MjSpec, mj_forward

from judo.tasks.cartpole import XML_PATH as CARTPOLE_XML_PATH
from judo.visualizers.model import ViserMjModel

# create a cartpole model
spec = MjSpec.from_file(str(CARTPOLE_XML_PATH))
model = spec.compile()
data = MjData(model)

# create the passive visualizer instance
server = viser.ViserServer()
viser_mj_model = ViserMjModel(server, spec)
breakpoint()  # open link - you should see the model here

# set a state
data.qpos[0] = 0.5
mj_forward(model, data)  # need to do this to update the body states

# update a state
viser_mj_model.set_data(data)  # the model should have updated!

Hopefully this is what you were asking about. Any suggestions for improving our docs on these lines is much appreciated!

Alternatively, if what you're asking about is new tooling to support these passive use-cases, let us know, and we can try to think of something that can work idiomatically in a way similar to the mujoco passive viewer. But, as-is, you could easily wrap your state updates in a for loop and use the minimal pattern above to hopefully achieve what you wanted.

[LaVieEstDure]

Hey @alberthli, thanks for responding! I actually did start fiddling with it yesterday and got it working in a passive visualiser usecase:

It was pretty simple as you point out and I got some custom functionality working! I can pass off the state to the trajectory optimiser, pass it a lambda that can generate trajectories and have it show the rollouts, control the rollout speed, etc.

I think it's very useful as a passive visualiser and has good potential for this usecase given how easy it was to add other UI elements. My thoughts on improvements on that front would be:

• I think the API for visualising traces is a bit more coupled to the visualiser node and realtime MPC than ideal. Sometimes you may want to visualise traces from things that aren't site sensors, or you may want the traces to be updated from various iterations of an iterative algorithm, and the not all usecases necessarily want to show elite vs non elite traces.
• One of the biggest challenges with sampling based trajopt is biasing sampling towards useful but diverse trajectories. For complex dynamic systems, sometimes a trace isn't enough to visualise those. It would be great to have some API to visualise multiple rollouts with multiple 3d models + some transparency.
• It might be worth providing tools for some extra visualisation functionality like visualising contact forces with arrows (exists in mj_passive) and visualising forces from actuators with arrows (e.g. thrusters/propellers, doesn't exist in mj_passive, it does changes colours though). Based on a quick look, looks like viser doesn't support adding arrows so maybe there's some scope for upstream work too.

I hope that at least some of that isn't going far too beyond the scope of judo. Would be keen to hear your thoughts

[alberthli]

Thanks for the detailed feedback. First, I have some high-level thoughts on scope.

One thing I think we're trying very hard to balance is having judo be feature-rich but also "appropriately-scoped". That is, our goal isn't to develop a package that does everything, but rather, to have one that is performant but simple enough for people to hack for sampling-based MPC applications. Your offline trajectory visualization use-case is a very well-motivated one that is enabled by the tooling in judo, but also doesn't necessarily feel like a "core" feature (at least until we expand the tooling significantly towards the data collection direction, which we have discussed a lot). I think maybe we could consider putting together an auxiliary judo-extras package that also lives in this repo but can be optionally installed with some of these additional features as a sibling to judo, and we can develop those features out on a slower timeline, maybe with community support. Thoughts @pculbertson @slecleach?

Regarding your specific bullet points:

• I totally agree that the trace visualization is a very specific type of visualization that doesn't always capture information well (and sometimes you want to register traces in a more ad hoc way). Do you have a suggestion for an interface for trace visualization that doesn't depend on sites?
• I don't think it would be too challenging to visualize multiple models along the traj, with the caveat that we've had trouble getting nice transparent visualizations in viser (maybe we're just not doing it correctly). There may also be concerns about speed if we're visualizing so many meshes at once.
• I agree that arrow visualizations would be really useful. However, it feels less important/impactful than all the other stuff discussed in this issue.

I've opened three sub-issues to keep track of your requests. Let us know if you would be willing to help develop some of these!

[pculbertson]

Thanks so much for the comments @LaVieEstDure!

Totally agree with Albert here - we're hoping to keep the core feature set quite focused, but balance this against making everything hackable. I like the judo-extras idea- another is that we could add lots of Viser-style usage examples to our docs (Brent had like 20 different examples of how to combine the Viser tools to make different apps), while keeping our core app focused on the async case.

Re: multi-mesh visualization, I'd anticipate some of the same problems Albert mentioned - for transparent objects, Three.js has trouble maintaining a consistent draw order. But generalizing the trace/force visualization seem like great ideas we could put on main, since these could be toggled with buttons like in the real MuJoCo viewer also.