Impulse vs XPBD; and assorted thoughts about brax HOT 34 OPEN

It seems like almost everything in the physics submodule is indeed specific to the physics simulator provided out of the box with brax currently; except for the math.py which provides more generic functionality. I suppose colliders.py could also be spit into collision detection and collision resolution for reusability.

In the end all that we want the physics module to do is provide a function next_state = physics.step(state, action), right? Any code outside does not particularly care about the layout of state or action I figure. The only other element of coupling is in the env specification; that currently is tailored for hierarchical rigid body formulations; I suppose it could quite easily be generalized to XPBD needs by adding an optional joint compliance field; but if you are thinking about adding a pure mass-spring physics backend as well, itd probably need its own environment spec format as well. Its hard to foresee all the potential points of coupling without having tried; think I will just start hacking away at my XPBD port first without much regard for generality and see where that takes me.

from brax.

Another interesting exercise in generality along the same lines would be to add support for a 2d backend as well. Having that also tends to be nice for debugability and testability.

from brax.

Hi Eelco! All of your suggestions/observations are super interesting! As a rough top-of-line response before I get into specifics, we'd be more than happy to accept PRs that provide new functionality, so we definitely encourage you to fork and hack away.

In rough order:

Re XPBD: We had not thought about XPBD at all, though we have vague plans of trying reduced coordinate methods like Featherstone at some point. I agree that it would be pretty slick to be able to toggle this kind of functionality on/off as desired. The integrator code actually used to be a bit more modular when I was playing with Runge-Kutta vs. Symplectic methods, so I'd be happy to open that back up again.

Re walking policies via gradients: To be more specific, we could absolutely replicate the 2D difftaichi walker env results with gradients, and this was actually one of the very first thing we did waaaaaay early in the project (though brax has undergone many changes since then). I was more referring to getting locomotive gaits with the 3D environments we included in the initial release. The issue is more a silly problem of optimization than it is of some fundamental issue with Brax: gradients through long trajectories are 1) super expensive and 2) super chaotic. For 1) we start OOMing when trying to take a gradient through a large batch of, say, Ants over 1000 steps. The "right" way to solve this would be to chunk up those 1000-step trajectories into random (say) 50-step sub-trajectories and then estimate the gradient via these truncated gradients (similar to Truncated Backprop Through Time). This simultaneously solves 2) because the gradient variance of shorter trajectories can be way smaller (though biased).

All that said, we definitely wouldn't say no to a method that's more careful about contact physics :)

Re reorganization for modularity: indeed, the current implementations are fairly specific to impulsive updates and hierarchical rigid bodies. One high-level design decision we'd like to preserve is easy vectorization of updates (i.e., vmap-ability of physics primitives), but beyond that, yes next_state = physics.step(state, action) is fairly generically workable.

Re 2d testing: We support this actually via the frozen field in the proto. The reacher environment is probably the simplest example of this to tinker with. I'm actually tinkering with a modified version of this env to run some differentiability experiments for our NeurIPS reviewers that you might find useful. I'll drop a comment here when it's in a more usable state :)

Thanks again for the interest! Definitely happy to collaborate on some cool new features!

from brax.

Hi Eelco - thanks so much for sharing these thoughts. Super cool.

+1 to Daniel's suggestion of sharing his colab. We have a colab that shows the best motivating example we have right now of using the differentiable properties of Brax. It shows that using analytical policy gradients, we can improve sample efficiency of solving the Reacher task by more than an order of magnitude.

That's great, but there are also some limitations:

• We suspect that sample efficiency can be even better.
• We don't have any results yet with environments that have contact.
• Gradients are surprisingly slow to calculate, and wall clock time is actually slower for the gradient based approach compared to PPO
• jit times of grad(env.step) can be surprisingly high

We are very interested in proofs of concept that address any of those issues - this XPBD approach sounds very interesting! If you can show some improvements on any of those bullet points above, we'd love to see them. And to Daniel's point, we can figure out the right way to incorporate those changes, there's a few ways to slice the codebase to make it more modular.

from brax.

Hi Daniel; Thanks for a cool project, and thanks for taking the time to reply!

Re XPBD: We had not thought about XPBD at all, though we have vague plans of trying reduced coordinate methods like Featherstone at some point. I agree that it would be pretty slick to be able to toggle this kind of functionality on/off as desired. The integrator code actually used to be a bit more modular when I was playing with Runge-Kutta vs. Symplectic methods, so I'd be happy to open that back up again.

Alright!

Re walking policies via gradients: To be more specific, we could absolutely replicate the 2D difftaichi walker env results with gradients, and this was actually one of the very first thing we did waaaaaay early in the project (though brax has undergone many changes since then). I was more referring to getting locomotive gaits with the 3D environments we included in the initial release. The issue is more a silly problem of optimization than it is of some fundamental issue with Brax: gradients through long trajectories are 1) super expensive and 2) super chaotic. For 1) we start OOMing when trying to take a gradient through a large batch of, say, Ants over 1000 steps. The "right" way to solve this would be to chunk up those 1000-step trajectories into random (say) 50-step sub-trajectories and then estimate the gradient via these truncated gradients (similar to Truncated Backprop Through Time). This simultaneously solves 2) because the gradient variance of shorter trajectories can be way smaller (though biased).

I dont have any hands on experience with these long rollouts myself yet; other than using checkpointing in large neural nets, which is a technique that might be applicable here too. On a related note it seems to me a priori that forward mode diff would be preferable for the physics step; this long compute graph with all these relatively tiny pointwise ops rather than nice chunky conv or mm ops seems like itd be a memory-allocation-and-fragmentation nightmare to run in backwards mode. Plus the benefits of backwards mode wouldnt really materialize anyway for all these pointwise ops i think. But if you run it together with a big NN controller you might want some hybrid composition of forward and backward gradients. Eh not talking from experience here so ill shut up.

Eventually I do hope to be able to do fairly long rollouts, as I intend to experiment with controllers that are able to observe a substantial state history backwards in time (you cant really observe if you are slipping from a single state using realistic sensors, for instance). How chaotic the gradients are is problem-specific I suppose; my application domain is single-track wheeled vehicles which I suppose are somewhat inherently more differentiable domains than intermittent-contact walking type problems. For a walker the dominant relevant time period is usually a single walking cycle; but for my wheeled vehicle I want my controller to be able to reason about how to initiate a turn while anticipating perhaps losing traction halfway, and how to pull out of that slip again. The natural cycle to reason about here is perhaps on the order of tens of seconds. Having some method within brax to do more-or-less arbitrarily long rollouts, if it makes sense within the domain, does seem useful to me. Doing it in under a minute is nice, but ill wait a day if I have to.

All that said, we definitely wouldn't say no to a method that's more careful about contact physics :)

I dont know that XPDB is more careful about contact physics per se; its in essence an additional integration of the impulse paradigm; with impulses, instead of the unit of account being forces, we reason about the effect of their once-time-integrated counterparts; and with (X)PBD, energy-and-momentum conserving displacements are the unit of account. Thats great for the stiffness of the contacts you can manage with an explicit integrator, but what that means for its suitability in a differentiable simulator is an entirely open question to me. When I said it looks suitable for a differentiable simulator I suppose what I meant was it seems simple enough to implement with my novice jax skills (as opposed to the potential practical trouble I foresee in getting an implicit solver to be effectively differentiable; gradsim claims to do it but its more than I am willing to tackle in my spare time). Rereading the intro of the linked XPBD paper they actually mention the resurgence of penalty-force based contact due to their smooth trajectories in differentiable physics; kinda sounds like they imply XPBD is a step in the wrong direction in that regard? Though they dont explicitly say that. I hope to find out; if nothing else I can save other people the trouble.

Re reorganization for modularity: indeed, the current implementations are fairly specific to impulsive updates and hierarchical rigid bodies. One high-level design decision we'd like to preserve is easy vectorization of updates (i.e., vmap-ability of physics primitives), but beyond that, yes next_state = physics.step(state, action) is fairly generically workable.

For sure, the vectorizability seems essential to me as well. In broad strokes the XPBD implementation should be rather similar to the one you already have, and there should be no difference in vectorizability, that I foresee.

Re 2d testing: We support this actually via the frozen field in the proto. The reacher environment is probably the simplest example of this to tinker with. I'm actually tinkering with a modified version of this env to run some differentiability experiments for our NeurIPS reviewers that you might find useful. I'll drop a comment here when it's in a more usable state :)

Thanks!

Thanks again for the interest! Definitely happy to collaborate on some cool new features!

Glad to hear; im very excited about differentiable physics and modern compilers like JAX, and I think you guys have laid the foundation for a super nice project here. While I think difftaichi and gradsim are super cool projects too, im a little wary of making your own DSL without institutional support; old enough to have seen that get stuck a few times, and I think building this type of infrastructure on top of JAX, while perhaps a bit more limiting in some ways, is probably the approach more likely to catch on in the medium term for many applications.

from brax.

Hi Eelco - thanks so much for sharing these thoughts. Super cool.

Thanks for making an awesome package!

+1 to Daniel's suggestion of sharing his colab. We have a colab that shows the best motivating example we have right now of using the differentiable properties of Brax. It shows that using analytical policy gradients, we can improve sample efficiency of solving the Reacher task by more than an order of magnitude.

That's great, but there are also some limitations:

• We suspect that sample efficiency can be even better.
• We don't have any results yet with environments that have contact.
• Gradients are surprisingly slow to calculate, and wall clock time is actually slower for the gradient based approach compared to PPO
• jit times of grad(env.step) can be surprisingly high

We are very interested in proofs of concept that address any of those issues - this XPBD approach sounds very interesting! If you can show some improvements on any of those bullet points above, we'd love to see them. And to Daniel's point, we can figure out the right way to incorporate those changes, there's a few ways to slice the codebase to make it more modular.

Not sure XPBD will help much with any of these. Compilation times and runtime characteristic for a single timestep will be similar I imagine; though arguably one might take bigger timesteps for the same stiffness being simulated using XPBD. XPBD is unconditionally stable and can do infinitely stiff joints; the only failure mode is that you get more compliant joints than you ordered if you are not putting in sufficient timesteps to realize all that stiffness. But I do think the XPBD authors convincingly demonstrate its competitiveness in terms of stiffness-per-flop, and it should outcompete an impulse based method in that regard, by perhaps an order of magnitude or so, for typical scenarios.

But if its a good match for differentiable simulation, I would not dare say yet. Looking at the implementation of XPBD, I dont see any steps in particular that would cause (additional) trouble.

However, the differentiability of the trajectory of a single body under elastic collision seems subject to similar considerations that difftaichi explores. XPBD 'integrates out' the contact dynamics and resolves it immediately within the timestep. It aims for zero penetration at the end of the timestep, regardless of time-of-impact. So thats not going to work without modification for a billiards simulation I think. Rather than the reverse gradients observed in the difftaichi paper id expect a stepwise zero-gradient result. (local gradients suggest it doesnt matter if your foot gets there earlier; the outcome appears the same; so why bother changing upstream the trajectory of your foot?). I really dont know how all that is going to work out in practice or what fixes will present themselves along the way, so I figured I just take the plunge.

And moreover, I don't think I need to care for my wheeled vehicle simulation and its continuous contact anyway (happy to leave jumps for the future for now).

from brax.

Ah yes, we have definitely had to increase the number of timesteps to account for joint stiffness - see the substeps param in the config. Both humanoid and halfcheetah have them cranked up pretty high.

Very curious to see what you cook up!

from brax.

Was planning on spending a solid block of time working on this earlier in the year, but life decided otherwise. However, just wanted to let you know this is still very much on my mind; picking up smaller projects now to get my feet wet with brax/jax in what spare time I do have, but with the intent of getting back to this eventually.

from brax.

FYI we've moved to position-based dynamics as of v0.0.11: https://github.com/google/brax/releases/tag/v0.0.11

We're quite pleased with the improvements from using it!

from brax.

Ah awesome; ill go check it out, must be nice to get rid of stability worries. Is there a net efficiency gain, on typical benchmarks? Guess youd want to compare with 'equal effective stiffness'.. so it might not be easy to get a real apples-to-apples.

Anyway... see this issue is almost a year old. Time flies with a second child I guess. Good hope ill be able to contribute more actively in the coming year though!

from brax.

Some quick feedback; looking good! It appears in one of my test cases that the effective static friction has gone down? As far as I can tell from looking at some other examples it is still supposed to be specified in the same manner?

No support for joint/material compliance yet that I can see; but I suppose thats on the roadmap? Should be easy to add on the XPBD side as I understand it no? I guess most of the work would be in propagating those properties through the rest of the codebase.

from brax.

Cant say ive fully digested it yet; but one oddity I noticed is that the OneWayCollider is dealing with the friction coefficient in its _position_contact method, whereas the TwoWayCollider isnt. Cant quite make sense of that as an intentional thing? Also for the sake of code simplicity wouldnt it be better to leave out this distinction in the XPDB case? OneWay could just be twoway with a zero inverse-mass assigned to one of the bodies, right? I guess feeding in that zero as a compile time constant could lead to some runtime efficiency gains; but do they in practice, and does that warrant the additional code complexity?

from brax.

If I zero out the friction local var in OneWayCollider _position_contact, I seem to get unbreakable static friction. That strikes me as the opposite behavior from what it should be; both from what youd expect from a friction coefficient; and as I read the code:
static_mask = jp.where(jp.abs(dlambdat) < jp.abs(friction * dlambda), 1., 0.)
That should be all zeroes, if friction = 0

It seems to me that this is due to an interaction with the _velocity_contact method; if I zero the friction there but leave the position based one intact, I also get a behavior resembling high static friction. But with both enabled, they seem to partially cancel eachother; leading to seemingly rather low static friction overall. I guess this might have to do with how both of these updates tie in with the integration loop?

from brax.

Ah interesting, this looks like a bug. Investigating!

No support for joint/material compliance yet that I can see; but I suppose thats on the roadmap?

Yeah not yet, primarily because we don't actually have any systems in brax that need springy joints. Also, PBD-style compliance adds another inner loop to optimize the lagrange multiplier, which isn't great for perf. But it'll find its way in eventually, I imagine.

from brax.

In one of the recent PDB papers they argue that its best anyway to leave the iteration count of that loop at 1; and throw all compute at the outer loop. So not bothering to store the intermediate multipliers should still fly, if i understand correctly.

I also dont have an application for springy joints myself at the moment; the thing that was most on my mind was springy contacts; which I need for reasons of physical approximation (the difference between say a rubber foot and a true hard contact can matter quite a lot for a policy I think). But also compliant contact, or the ability to anneal their stiffness over training, might be very useful with the whole gradient propagation / time of impact thing I imagine.

Would love to make a PR for it but my employer still hasnt given clarity on it.

from brax.

Also one thing I dont really understand; whats up with the 'performs two physics substeps per collision update' in _pbd_step? Is that a manual unrolling of the solver loop, effectively? Ive also seen this refered to elsewhere in the code, in _velocity_contact. It feels a bit brittle and tightly coupled the way its implemented, and im not sure I really understand the logic. Does this come from some reference where its explained in more detail?

from brax.

Alright, elasticity should be behaving correctly now. There was a bug making it basically always be elasticity=0.

We will probably do as you point out and reduce this to a single collision method at some point.

Also one thing I dont really understand; whats up with the 'performs two physics substeps per collision update' in _pbd_step? Is that a manual unrolling of the solver loop, effectively? Ive also seen this refered to elsewhere in the code, in _velocity_contact. It feels a bit brittle and tightly coupled the way its implemented, and im not sure I really understand the logic. Does this come from some reference where its explained in more detail?

This is entirely because collision are more expensive to calculate using PBD, so I'm simply checking them at a lower frequency than the other updates. We agree that this isn't aesthetically the best, and want to simplify this logic. I wasn't comfortable dropping perf by half across the board, though, which is what happens with collisions checks every step.

from brax.

Alright, elasticity should be behaving correctly now. There was a bug making it basically always be elasticity=0.

I hadnt paid attention to the elasticity yet. Friction still seems funny to me in the same way (and I didnt notice any code changes pertaining to it).

Also one thing I dont really understand; whats up with the 'performs two physics substeps per collision update' in _pbd_step? Is that a manual unrolling of the solver loop, effectively? Ive also seen this refered to elsewhere in the code, in _velocity_contact. It feels a bit brittle and tightly coupled the way its implemented, and im not sure I really understand the logic. Does this come from some reference where its explained in more detail?

This is entirely because collision are more expensive to calculate using PBD, so I'm simply checking them at a lower frequency than the other updates. We agree that this isn't aesthetically the best, and want to simplify this logic. I wasn't comfortable dropping perf by half across the board, though, which is what happens with collisions checks every step.

Ah ok; I recall something from the original XPBD paper as well, where they run at least the broad phase collision detection at a lower frequency I think.

from brax.

Re friction---would you mind sharing your unit test? There was some weirdness in how multiple-contact normal forces were being averaged that should be a little better behaved now.

from brax.

I do not think it has to do with multiple contacts; it shows up for a single one. Ive been thinking about how to make this into a small reproducible test but its part of a tangle of experimental code right now... fairly sure the problem is with the current PBD implementation; in any case the behavior is very different from the legacy constraint handling.

from brax.

Here is how a spinning box thrown on the ground behaves for me; not sure that either pbd or legacy is without quircks; but they sure are different. I can seem to reproduce the apparent high static friction in this simple example; but its food for thought by itself in my opinion. It is as if in the PBD case the friction doesnt apply to the rotational DOFs somehow?

_SYSTEM_CONFIG = """
bodies {
  name: "Ground"
  colliders {
    plane {}
  }
  inertia { x: 1.0 y: 1.0 z: 1.0 }
  mass: 1
  frozen { all: true }
}

bodies {
  name: "Rect"
  colliders {
    box {
      halfsize { x: 0.5 y: 1.0 z: 2.0 }
    }
  }
  inertia { x: 4.0 y: 2.0 z: 1.0 }
  mass: 1
}


gravity { z: -10.0 }
angular_damping: 0.0
# dunno what this is yet
baumgarte_erp: 0.2

friction: 0.5

dt: 0.05
substeps: 50
dynamics_mode: "legacy_spring"
"""

from brax.

Hmm I guess my moments of inertia are artificially inflated in the above example, contributing to a sense of irrealism (and the y-z dimensions of the cube are swapped in rendering; thanks vispy, that makes total sense). Note that between the two animations, only the PBD flag is different though.

Fixing that brings up another host of issues though; wont bother to list them all here since many if not most are probably simply noob mistakes by me still; but I think itd be good to add some (visual) unit tests to the codebase involving simple setups along the lines above, to have a clear baseline when making changes to the code.

from brax.

To be more clear about my comment above; these two lines still appear inconsistent with one another, wrt to the appearance of the friction variable in only one of the two.

from brax.

Some additional findings; there are various ways to trigger instabilities in the PBD simulation. Lower inertia values tend to be less stable. Blocks of styrofoam with a metal ball inside might not be too common but they re not unphysical per se. I dont recall any PBD paper fessing up to such instabilities intrinsic to the method; but that doesn't mean they dont exist of course. I imagine its probably a bug though. Also, higher friction values around 1 or upwards tend to be less stable. Some of these instabilities are present in the legacy mode as well, though it tends to be more reasonable in its behavior. Aside from the thing above, which strikes me as an unrelated bug, I havnt found any leads yet as to what might be the cause of this.

from brax.

Looking around for problems (and not finding root causes thus far) another nit to pick:

This line strikes me as undesirable compared to the two-way equivalent, in the sense that the behavior can be expected to differ from a two-way collider in the limit to infinite mass.

Im also a bit puzzled by the appearance of the .3 multiplication constant appearing in the twoway collider. Giving that constant a name might help clarify the intent behind it; but either way it strikes me as another difference between the oneway and the twoway code path. EDIT: I suppose its meant as an underrelaxation for the jacobi iteration? Thats an interesting topic actually. vectorization kinda demands jacobi iteration for efficiency; but guaranteed stability without additional tweaks demands gaussian iteration; I think? I suppose some kind of red-black or similar graph coloring scheme might be a good middle ground; but it shouldnt be the cause of instability in the single-contact scenarios im looking at. Or thinking about it again... it might actually. I guess a cube and plane have up to 4 active shared contacts at the same time. These constraints are quite colinear so updating them together, without regard for the effect they have on one another, can render plain jacobi unstable I think? Dont think that has bearing on my present problem either since there is plenty of funkiness with only a single active contact.... but it might be another rabbit hole worth looking out for. Maybe the simplest fix there is to downweight the jacobi update by the number of constraints active between each pair of bodies, rather than a fixed constant?

from brax.

I think im on to something here; slapping a 0.5 underrelaxation on the oneway position_contact and velocity_contact, gives me stable behavior. Both seem to be required for a cube hitting the ground with two corners at the same time not to freak out. This seems to be the more fundamental fix; if i downtune my inertias again to get my steel-ball-in-block-of-styrofoam scenario, its now stable; and behaves physically as expected.

Thinking about this again; jacobi iteration is supposed to be unconditionally stable for problems of this type; but it does require division by the diagonal of the system matrix as a whole (which is kinda but not quite like the downweighting of each constraints by the number of constraints between that set of bodies which i mentioned above; dividing by the number of constraints between the pair of bodies should be the more conservative approach, strictly leading to smaller updates than a diagonal division; and the two should coincide when the constraints are completely colinear). I suppose that diagonal-division might be missed from the code as its implemented now? I cant readily identify it in the logic.

I know of some partial PBD reference implementations; but they all do GS iteration afaik. Either way im sure this jacobi situation is fixable in an elegant manner; but ive only seen it mentioned obliquely in PBD texts.

Also also someone who is still new to JAX and doesnt have a good intuition for it; I wonder how feasible GS in JAX would be; it depends on the degree to which one depends on parralelism over the number of constraints to get decent performance. For single rollouts at a time, im pretty sure this is going to have very suboptimal on-device performance. But is that the use case we should be concerned with? For large enough batches, I suppose one doesnt care.

from brax.

@cdfreeman-google @erikfrey Not sure spamming this issue is the best place for these comments by the way; if you have a more convenient place for this type of open end chat, let me know!

from brax.

Ah these guys do go into some detail about jacobi in PBD

https://www.researchgate.net/publication/316433229_Unified_Simulation_of_Rigid_and_Flexible_Bodies_Using_Position_Based_Dynamics

'The actual spectral radius ρ(N) does not need to be computed, as it can approximated from the count of incident constraints to a body'; I think thats basically saying the same as what I said above; but not quite?

Dividing each update by the count of incident constraints does not sound ideal to me; as it breaks momentum conservation of the update. Doesnt matter if you iterate till convergence, but I think its preferable to stay on the manifold of conservation-respecting updates. Dividing the per-constraint stepsize does preserve momentum conservation; but what to divide it by? I suppose the number of constraints between those bodies isnt a failsafe solution; if you have a cube sitting on top of four independent cubes with each concern, you have the same issue as for a block sitting on a plane. So damping each constraint update by the max incident constraint count of any of the bodies it affects might be an option as well; though it will give more cautious stepsizes in general I suppose. Many constraints are not very colinear at all and wed be brute-force damping their relaxation as well.

Eh anyway here is me complicating things again; I suppose just following the logic of that paper first would be a good start; they contain some seemingly interesting references on this topic as well that i havnt checked out yet.

from brax.

I tried this 0.5 jacobian damping in the oneway collider in my non-synthetic test case, and there it also seems to get rid of the oddities I was observing earlier; at least at first glance; certainly more plausible behavior.

from brax.

Haha no this is great---thank you tremendously for looking into this! I can confirm that lowering the strength of the collision interaction term stabilizes things for me as well (I've actually been doing that in a branch this week with a particularly low inertia system). I'll expose this value in the config in the next update.

I'll take a look at that ref--I agree breaking momentum conservation is not great. (sorry for the lack of comms on my side!!)

from brax.

Hey no problem; ive been meaning to dig into this for a year now, ive got no right to be pushy :). Either way if you have a channel for extended discussion id be interested in joining in.

Adding a configurable value might indeed be a decent quick fix to help things along; but I think there shouldnt be a single value that is unconditionally stable; find a way of slapping more colinear constraints together and your problem will be back. But yeah I think it would be helpful until a more principled solution is found and validated to work.

Not sure if the breaking of momentum conservation would be an issue in practice. The cited paper doesnt mention it; and again it shouldnt matter if you iterate jacobi to convergence. But given the prevailing modern PBD wisdom of 'a single solver iteration is optimal; just run your outer loop more frequently', I could imagine this becoming a problem. All these PBD formulas can be viewed as being derived from considering what is action-minimizing, momentum-preserving step to to bodies involved, to satisfy each constraint in isolation at the end of the timestep. If those computed steps arnt applied equally to all bodies involved in the constraint, seems to me that could quickly lead to real world observable consequences. But I could be wrong; and that particular scheme of just counting all constraint incident to each body should be easy to implement.

Ah; only just noticed this line!

from brax.

One more thing to consider:

I just realized that different joint types and different contact types are already handled in a serial manner; at least reading the code as-if it were python; I suppose the JAX compiler could dispatch the different terms in those sums for the joint-type and collider-types in parralel; but even if it did (no clue if it would?) those kernels themselves would only gain minimal parallelism in many scenarios, unless there is actually a large number of constraints of a given type (which for my use cases will not be the case). I guess unless there are two collider or joint types that both have a substantial count of them, parralelizing their invocation wont do much.

I suspect that at least visiting in a GS-fashion per constraint type would be a win for robustness; still have to worry about the stability of the set of constraints of one type updated in parralel; but if you end up needing to reduce their stepsize to respect stability within that jacobi-updated block, at least those conservative stepsizes will not spill over unnecessarily, and we are using fresh information as much as possible.

from brax.

One more thing to reflect on; I see that the actuators are still implemented in terms of impulses.

As I understand, many hardware implementations have a PD feedback loop running at high frequency, and the high level controller is trained to update the setpoint of this PD controller at a lower frequency. For very stiff controller parameters, trying to simulate such a system using an impulse description could become unstable just like the other type of constraints; and itd be preferable to implement the motor-controller system as a PBD constraint as well. Seems like that should be simple to implement; just need to find the right mapping from motor-controller system parameters to the compliance and damping of the PBD constraint; otherwise the actuator implementation should follow the joint implementation pretty closely?

Not sure if itd be easy to model an integral-term controller in such a manner as well; seems like a chore having to drag that additional state of the controller around in your code; but perhaps there is an elegant trick. But I also doubt it adds much; in most real world scenarios the higher level controller should probably be able to learn to compensate for the absence of a windup term by overshooting the setpoint, since steady-state errors are an effect only relevant over longer timescales.

Dont think that its a super urgent concern of mine, but in the long term itd be nice if the system was general enough that itd be intuitive to set up an entire drivetrain, with motors having physically correct inertias and gyroscopic effects, belt drives with realistic springiness, etc, so the whole electro-mechanical system can be simulated.

EDIT: speaking of belt drives; they are also interesting from a code structure POV I suppose. If you want them to be slip-free timing belts (or gears), you also need to maintain some kind of belt-stretch auxillary variable that you pass from timestep to timestep. Otherwise the stretch will be 'forgotten' at the start of each step, I suppose, and a belt under continuous load would drift over time. One question that only occurs to me now for the first time... is it actually feaible to directly simulate objects like high-RPM motor-rotors? Generally for collision resolution youd want to resolve each revolution in many steps; but for a rotationally symmetric object that shouldnt really collide anyway, its fine I suppose. But the logic of PBD also requires at minimum two timesteps per revolution, or otherwise the basic integrator just falls apart, with angular momentum suddenly reversing if you go past half a rev per timestep. On top of that, I think getting physically correct gyroscopic forces also depends on actually resolving the actual rotation of the object with reasonable resolution. Then again, seems most env currently run at about 1000 substeps per second; if 10 steps per revolution would suffice, thats 100revs/sec, or 6000rpm. So perhaps 'direct numerical simulation' of motor components actually isnt completely outside the realm of reason. Thatd be great; because having to make some simplified/lumped model or such components and having to figure out a way to correctly couple those to the rest of the simulation sure would be less elegant.

from brax.

An additional note: the way the joint damping is applied in _pbd_step seems to deviate from the XPBD paper; its explicitly integrated rather than handled together with the other friction calculations. Dont know if the implementation intends to follow one paper or another explicitly; but if its an intentional difference would be good to call it out in a docstring I think.

Note that the paper I linked above (see eq 20 and related section) does handle joint damping in a yet different, and in my opinion superior, more PBD consistent manner. Rather than adding it as an acceleration, or velocity correction term, its added as an extra term in the PBD position-step solving logic. In this way, any competition between stiff spring forces and high damping would be correctly balanced out on the sub-timestep scale.

Doesnt really add any meaningfull code complexity that I can see, and it would also make this damping unconditionally stable, regardless of the damping coefficient. Not sure it is immediately relevant to any envs that are currently implemented; but if you wanted to simulate something super damped like a hydraulic cylinder, or a stiff rubber element, itd probably be welcome.

EDIT: An attempt to gather some additional input.

EDIT2: thinking about it a little; I imagine the XPBD route should indeed also be unconditionally stable. I suppose the main difference would should up in overdamped scenarios. Given a joint out of equilibrium, using the XPBD approach it would first have its position error corrected, then it would get its velocity set, and then the velocity would be zeroed out again in the damping step. Using the method linked above however, an overdamped joint would correctly balance the damping and joint stiffness; and if the damping dominates the joint should creep up to its equilibrium position over multiple timesteps. For underdamped systems id imagine they converge to the same behavior. But I think the rayleigh damping is also attractive because it should generalize in a nice manner to arbitrary constraint types.

from brax.