newton.solvers.SolverBase#

class newton.solvers.SolverBase(model, *, collision_pipeline=None, collision_frequency=None, collision_frequency_type=None)[source]#

Bases: object

Generic base class for solvers.

The implementation provides helper kernels to integrate rigid bodies and particles. Concrete solver back-ends should derive from this class and override step() as well as notify_model_changed() where necessary.

class CollisionFrequencyType(*values)#

Bases: IntEnum

When, inside a step(), a solver-owned collision pipeline runs detection.

The frequency number in collision_frequency applies only to ITERATIONS; the other members ignore it. Skipping detection across steps carries no hidden solver state — set a slot to NONE between steps via set_collision_frequency().

AUTO = 4#: Solver-specific default.

ITERATIONS = 3#: Before initialization, then immediately before iterations k, 2k, and so on.

NONE = 0#: Never detect; the user may run detection externally into contacts.

PRE_INIT = 1#: Once per step, before solver initialization.

PRE_POST_INIT = 2#: Before and after solver initialization (one detection each).

class CollisionSlot(*values)#

Bases: IntEnum

Collision-detection categories scheduled by a solver.

RIGID = 0#: Rigid-rigid and particle-shape collision detection.

SOFT_SELF_CONTACT = 1#: Triangle-mesh soft self-contact detection.

class ObservableKind#

Bases: object

Inheritable names for standard and solver-specific observable kinds.

Access kinds through solver.ObservableKind and compose requests as sets. This is an ordinary namespace class of string constants: solvers may subclass it and declare new entries, such as BODY_TEMPERATURE = "body_temperature". Inherited names remain available. Literal strings and these constants are interchangeable in requests. Kind strings must be unique within a solver’s observable declarations.

Subclasses are checked when defined. Every public attribute must be a nonempty string. Changing an inherited name’s value or assigning one string to different names raises an error, including conflicts between multiple parents. Redeclaring the same name with the same value is allowed. Attributes beginning with an underscore are excluded from these checks.

The available names do not imply solver support. Query SolverBase.supported_observables for the effective capabilities of a configured solver. Use that set for iteration or requesting all supported observables; the namespace itself is not iterable.

See Solver Observables for request and extension examples.

Experimental

The solver observable API may change without prior notice.

BODY_PARENT_F = 'body_parent_f'#: Incoming parent-joint wrenches on rigid bodies.

BODY_QDD = 'body_qdd'#: Rigid-body spatial accelerations.

CONTACT_F = 'contact_f'#: Spatial contact forces aligned with a Contacts container.

class Observables#

Bases: object

Arrays populated by a solver in addition to the simulation state.

Allocate a container once with SolverBase.observables(), then pass that same container as observables= to the allocating solver’s SolverBase.step() on each step that needs diagnostics. The solver writes requested results into the existing arrays during each such call; allocation alone does not register them for automatic updates.

Solver implementations can derive from this class to add solver-specific arrays while retaining the standard Newton observables. Decorate derived containers with @dataclass(eq=False) and declare arrays with field(). Identity equality keeps containers usable as sources in ArticulationView caches.

See Solver Observables for allocation, selection, and custom field examples.

Experimental

The solver observable API may change while additional solvers and observable categories are migrated to it.

static field(*, dtype, frequency, kind=None)#

Declare an optional observable array on a dataclass container.

See Solver Observables for custom observable declarations.

Parameters:

dtype (type) – Warp element type, such as wp.float32 or wp.vec3.
frequency (AttributeFrequency | str) – Row domain, not update cadence. Custom string frequencies use model count and articulation-ownership metadata.
kind (str | None) – Nonempty string identifying this observable. If omitted or None, use the Python field name.

Returns:

A dataclass field defaulting to None. No array is allocated until SolverBase.observables() requests it.

Raises:

TypeError – If an explicit kind is not a nonempty string or the frequency is not an attribute frequency or nonempty string.

Return type:

Any

Experimental

Experimental feature. API, behavior, defaults, and supported use cases may change without prior notice.

__init__()#

get_attribute_frequency(name)#

Return an array’s row domain, including inherited declarations.

Parameters:

name (str) – Observable array field name.

Raises:

KeyError – If no frequency is declared for the field.
TypeError – If the container does not use an identity-based dataclass.

is_requested(kind)#

Return whether this container requests an observable on the current call.

This checks the request, not whether its values are fresh. It also returns True for requested zero-length arrays and during allocation.

Parameters:: kind (str) – A standard or solver-specific observable name.

select(kinds)#

Return a reusable subset sharing this container’s allocated arrays.

The result has the same concrete type, solver owner, and selected array objects, including gradients. Fields omitted from the selection are None in the result; this container is not modified. Stepping a subset therefore leaves the source’s omitted arrays unchanged.

Create selections before graph capture and reuse them across substeps. Only Python containers are created; no array allocation or copying is performed. Selections share the source’s contact-storage binding.

Derived containers with nested observable containers should override this method, call super(), and select their children in the result. Other solver-specific metadata is shallow-copied.

See Solver Observables for selection examples.

Parameters:

kinds (Iterable[str]) – Subset of kinds to request. An empty set requests no observables. Selecting a selection can only narrow it.

Returns:

A same-type container referencing the selected arrays.

Raises:

TypeError – If kinds is not a collection of nonempty strings.
ValueError – If this container was not allocated by a solver or a kind is not requested by this container.

Return type:

Observables

body_parent_f: wp.array[wp.spatial_vectorf] | None = None#: Incoming parent-joint wrenches [N, N·m], shape (body_count,).

body_qdd: wp.array[wp.spatial_vectorf] | None = None#: Rigid-body accelerations [m/s², rad/s²], shape (body_count,).

contact_f: wp.array[wp.spatial_vectorf] | None = None#: Contact forces [N, N·m], shape (rigid_contact_max + soft_contact_max,).

property contacts: Contacts | None#: Contact storage supplied at allocation, or None without contact requests.

property kinds: frozenset[str]#: Observable kinds requested by this container or selected subset.

property model: Model | None#: Model whose indexing this container uses, or None before allocation.

classmethod register_custom_attributes(builder)#

Register custom attributes for the solver.

Parameters:: builder (ModelBuilder) – The model builder to register the custom attributes to.

__init__(model, *, collision_pipeline=None, collision_frequency=None, collision_frequency_type=None)#

Initialize common solver state and optional collision scheduling.

Parameters:

model (Model) – Simulation model integrated by the solver.
collision_pipeline (CollisionPipeline | None) – Collision pipeline owned and driven by the solver. The pipeline must use model, and the concrete solver must set supports_collision_pipeline.
collision_frequency (Mapping[CollisionSlot, int] | None) – Per-slot iteration frequencies. Values must be at least one and are used only for slots scheduled with CollisionFrequencyType.ITERATIONS. Unspecified slots retain their defaults.
collision_frequency_type (Mapping[CollisionSlot, CollisionFrequencyType] | None) – Per-slot detection points. Unspecified slots retain their defaults.

integrate_bodies(model, state_in, state_out, dt, angular_damping=0.0)#

Integrate the rigid bodies of the model.

Parameters:

model (Model) – The model to integrate.
state_in (newton.State) – The input state.
state_out (newton.State) – The output state.
dt (float) – The time step (typically in seconds).
angular_damping (float) – The angular damping factor. Defaults to 0.0.

integrate_particles(model, state_in, state_out, dt)#

Integrate the particles of the model.

Parameters:

model (Model) – The model to integrate.
state_in (newton.State) – The input state.
state_out (newton.State) – The output state.
dt (float) – The time step (typically in seconds).

notify_model_changed(flags)#

Notify the solver that parts of the Model were modified.

The flags argument is a bit-mask composed of the ModelFlags enums or custom int bits. Each flag represents a category of model data that may have been updated after the solver was created. Passing the appropriate combination of flags enables a solver implementation to refresh its internal buffers without having to recreate the whole solver object. Solver implementations must handle each relevant narrow DOF flag: checking JOINT_DOF_PROPERTIES alone does not detect them. The broad flag covers force, inertial, and reference properties; friction/damping belongs to the force category. Solvers may ignore categories they do not cache or support, but should preserve unrelated pending edits except for documented backend dependencies. In MuJoCo, notifications that recompute constants also republish the current tendon limit modes, gains, raw solref values, and ranges, even without TENDON_PROPERTIES. Force-only joint notifications leave these tendon parameters unchanged. Valid flags are:

ModelFlags.JOINT_PROPERTIES: Joint transforms, axes, or coordinates have changed.
ModelFlags.JOINT_DOF_PROPERTIES: Full joint DOF update, including force, armature, and reference-pose properties (legacy behavior).
ModelFlags.JOINT_DOF_FORCE_PROPERTIES: Joint friction, damping, target gains/modes, effort limits, passive stiffness, or limit coefficients/bounds have changed.
ModelFlags.JOINT_DOF_INERTIAL_PROPERTIES: Joint armature has changed.
ModelFlags.JOINT_REFERENCE_POSE_PROPERTIES: Joint reference poses or spring references have changed.
ModelFlags.BODY_PROPERTIES: Rigid-body pose or velocity buffers have changed.
ModelFlags.BODY_INERTIAL_PROPERTIES: Rigid-body mass or inertia tensors have changed.
ModelFlags.SHAPE_PROPERTIES: Shape transforms or geometry have changed.
ModelFlags.MODEL_PROPERTIES: Model global properties (e.g., gravity) have changed.
ModelFlags.CONSTRAINT_PROPERTIES: Constraint definitions, coefficients, or enable flags have changed.
ModelFlags.TENDON_PROPERTIES: Tendon stiffness or related tendon properties have changed.
ModelFlags.ACTUATOR_PROPERTIES: Actuator gains, biases, limits, or force properties have changed.

Parameters:: flags (ModelFlags | int) – Bit-mask of ModelFlags or custom int bits indicating which model properties changed.

observables(kinds, *, contacts=None, requires_grad=None)#

Allocate reusable arrays for requested solver observables.

See Solver Observables for allocation and usage examples.

A container is owned by the solver that allocates it and can be passed to that solver’s step() method on every time step. Derived solvers add custom kinds to SUPPORTED_OBSERVABLES and declare fields on a dataclass derived from SolverBase.Observables. The default allocator handles both standard and custom fields. Override this factory and delegate to super() when additional solver-specific initialization is needed.

All requested arrays are allocated before this method returns; unrequested fields remain None. Contact arrays use the model’s resolved rigid and soft capacities, not the live contact count. Allocate before graph capture and pass the same Contacts to step() and consumers.

Parameters:

kinds (Iterable[str]) – Required set or other iterable of observable names. Entries from ObservableKind and literal strings are interchangeable. Pass supported_observables to allocate all observables for this instance. An empty collection allocates no arrays.
contacts (Contacts | None) – Contact storage to bind to contact-indexed observables. Required when any requested field has contact frequency, including zero-capacity fields. Its device and rigid/soft capacities must match the model. Ignored when no contact-indexed fields are requested.
requires_grad (bool | None) – Whether allocated arrays require gradients. If None, use the model’s setting.

Returns:

A solver-owned observable container with requested arrays allocated.

Raises:

TypeError – If a request is not a collection of nonempty strings or the configured observable type does not derive from SolverBase.Observables, or required contacts are not a Contacts instance.
ValueError – If this solver does not support a requested observable or its container does not declare a unique field for the kind, or required contacts are missing or incompatible.
RuntimeError – If contact-indexed observables are requested before constructing CollisionPipeline for the model.

Return type:

Observables

Experimental

The solver observable API may change while additional solvers and observable categories are migrated to it.

reset(state, world_mask=None, flags=None)#

Reset the solver internal state data.

Modifies the given state in place. Derived solvers override this to reset solver-specific internal buffers or custom state attributes when environments are reset (e.g. during RL training).

The default implementation is a no-op so solvers that do not require special reset logic need not override this method.

Parameters:

state (newton.State) – The simulation state to reset (modified in place).
world_mask (wp.array[wp.bool] | None) – Optional boolean mask of shape (world_count + 1,) specifying which worlds to reset. Entries before the last select local worlds by index, and the final entry selects global entities whose world is -1. If None, all local and global entities are reset.
flags (StateFlags | int | None) – Optional StateFlags or int bitmask controlling which state attributes need to be reset. If None, all state attributes are reset.

set_collision_frequency(*, collision_frequency=None, collision_frequency_type=None)#

Change the detection schedule; takes effect at the next step().

The solver keeps no hidden cross-step scheduling state, so detecting every N steps is expressed by toggling a slot between CollisionFrequencyType.NONE and an active type from the calling loop. None keeps the corresponding current setting. Recapture an existing CUDA graph after changing the schedule.

Parameters:

collision_frequency (Mapping[CollisionSlot, int] | None) – Frequency numbers keyed by CollisionSlot; used only by ITERATIONS slots (before iterations k, 2k, and so on) and must be at least one.
collision_frequency_type (Mapping[CollisionSlot, CollisionFrequencyType] | None) – Detection points keyed by CollisionSlot.

step(state_in, state_out, control, contacts, dt, *, observables=None)#

Simulate the model for a given time step using the given control input.

Parameters:

state_in (newton.State) – The input state.
state_out (newton.State) – The output state.
control (Control | None) – The control input. Defaults to None which means the control values from the Model are used.
contacts (Contacts | None) – The contact information.
dt (float) – The time step (typically in seconds).
observables (Observables | None) – Optional solver observable arrays allocated by observables().

update_contacts(contacts, state=None)#

Update a legacy Contacts object with forces from the solver state.

Deprecated since version 1.7: Request ObservableKind.CONTACT_F using observables() and pass the resulting container to step() instead.

Parameters:

contacts (Contacts) – The object to update from the solver state.
state (State | None) – Optional simulation state, used by some solvers.

validate_observables(observables, contacts=None)#

Validate ownership and contact storage before a custom solver step.

Call this at the start of step(), before launching work or writing observable arrays. Passing None is a no-op. This checks only ownership and contact identity; layout validation happens during allocation.

See Solver Observables for custom solver integration.

Parameters:

observables (Observables | None) – Optional container allocated by this solver.
contacts (Contacts | None) – The instance supplied at allocation, required for contact-indexed observables.

Raises:

TypeError – If the container has an incompatible type.
ValueError – If the owner or contact storage is incompatible.

Experimental

Experimental feature. API, behavior, defaults, and supported use cases may change without prior notice.

SUPPORTED_OBSERVABLES: ClassVar[frozenset[str]] = frozenset({})#

Names of observables supported by this solver class.

Entries are string identifiers, such as ObservableKind.BODY_QDD. Query supported_observables for the capabilities of an instance.

property collision_frequency: dict[CollisionSlot, int]#: Per-slot detection frequency numbers as a read-only copy.

property collision_frequency_type: dict[CollisionSlot, CollisionFrequencyType]#: Per-slot CollisionFrequencyType values as a read-only copy.

collision_pipeline#: The solver-owned collision pipeline, or None when detection is driven externally.

property contacts: Contacts | None#

The solver-owned contacts buffer, or None when no pipeline is owned.

Unlike Model.contacts(), this property does not allocate; it returns the buffer created from the owned pipeline at construction. With a slot set to CollisionFrequencyType.NONE the user may fill this buffer externally, e.g. pipeline.collide(state, solver.contacts).

property device: Device#

Get the device used by the solver.

Returns:: The device used by the solver.
Return type:: wp.Device

property supported_observables: frozenset[str]#

Names accepted by observables() on this instance.

Defaults to SUPPORTED_OBSERVABLES. Solvers may narrow the class declaration for backend limitations, or derive capabilities from child solvers. Applications should query this property.

See Solver Observables for available observables and backend limitations.

supports_collision_pipeline: bool = False#

Whether this solver can own a CollisionPipeline and drive detection itself.

Currently only SolverVBD opts in; passing collision_pipeline to any other solver raises ValueError (drive detection externally instead).