dFlex is a physics engine for Python. It is written entirely in PyTorch and supports reverse mode differentiation w.r.t. to any simulation inputs.
It includes a USD-based visualization library (dflex.render
), which can generate time-sampled USD files, or update an existing stage on-the-fly.
- Python 3.6
- PyTorch 1.4.0 or higher
- Pixar USD lib (for visualization)
Pre-built USD Python libraries can be downloaded from https://developer.nvidia.com/usd, once they are downloaded you should follow the instructions to add them to your PYTHONPATH environment variable.
By default dFlex uses the built-in PyTorch cpp-extensions mechanism to compile auto-generated simulation kernels.
- Windows users should ensure they have Visual Studio 2019 installed
To use the engine you can import first the simulation module:
import dflex.sim
To build physical models there is a helper class available in dflex.sim.ModelBuilder
. This can be used to create models programmatically from Python. For example, to create a chain of particles:
builder = dflex.sim.ModelBuilder()
# anchor point (zero mass)
builder.add_particle((0, 1.0, 0.0), (0.0, 0.0, 0.0), 0.0)
# build chain
for i in range(1,10):
builder.add_particle((i, 1.0, 0.0), (0.0, 0.0, 0.0), 1.0)
builder.add_spring(i-1, i, 1.e 3, 0.0, 0)
# add ground plane
builder.add_shape_plane((0.0, 1.0, 0.0, 0.0), 0)
Once you have built your model you must convert it to a finalized PyTorch simulation data structure using finalize()
:
model = builder.finalize('cpu')
The model object represents static (non-time varying) data such as constraints, collision shapes, etc. The model is stored in PyTorch tensors, allowing differentiation with respect to both model and state.
To advance the simulation forward in time (forward dynamics), we use an integrator
object. dFlex currently offers semi-implicit and fully implicit (planned), via. the dflex.sim.ExplicitIntegrator
, and dflex.sim.ImplicitIntegrator
classes as follows:
sim_dt = 1.0/60.0
sim_steps = 100
integrator = dflex.sim.ExplicitIntegrator()
for i in range(0, sim_steps):
state = integrator.forward(model, state, sim_dt)
To visualize the scene dFlex supports a USD-based update via. the dflex.render.UsdRenderer
class. To create a renderer you must first create the USD stage, and the physical model.
import dflex.render
stage = Usd.Stage.CreateNew("test.usda")
renderer = dflex.render.UsdRenderer(model, stage)
renderer.draw_points = True
renderer.draw_springs = True
renderer.draw_shapes = True
Each frame the renderer should be updated with the current model state and the current elapsed simulation time:
renderer.update(state, sim_time)
Miles Macklin ([email protected])