planner.dynamics.mujoco_dynamics#

Module Contents#

planner.dynamics.mujoco_dynamics.get_joint_dimensions(joint_ids: numpy.typing.ArrayLike, state_address: numpy.typing.ArrayLike, state_length: int) → torch.IntTensor#

Given a list of joint ids, and the list of addresses in the states for the joints. We return the dimensions in the state corresponding to the list of joint ids.

Parameters:

  • joint_ids (numpy.typing.ArrayLike) –

  • state_address (numpy.typing.ArrayLike) –

  • state_length (int) –

Return type:

torch.IntTensor

planner.dynamics.mujoco_dynamics.decompose_state_dimensions(model: mujoco.MjModel) → Tuple[torch.IntTensor, torch.IntTensor, torch.IntTensor, torch.IntTensor]#

Decompose the states indices in 4 groups, split between positions or velocities and actuated or not actuated.

Parameters:

model (mujoco.MjModel) –

Return type:

Tuple[torch.IntTensor, torch.IntTensor, torch.IntTensor, torch.IntTensor]

class planner.dynamics.mujoco_dynamics.MujocoPlant(params: jacta.planner.core.parameter_container.ParameterContainer)#

Bases: jacta.planner.dynamics.simulator_plant.SimulatorPlant

Parameters:

params (jacta.planner.core.parameter_container.ParameterContainer) –

reset() → None#
Return type:

None

initialize(params: jacta.planner.core.parameter_container.ParameterContainer) → None#
Parameters:

params (jacta.planner.core.parameter_container.ParameterContainer) –

Return type:

None

dynamics(states: torch.FloatTensor, actions: torch.FloatTensor, action_time_step: float) → Tuple[torch.FloatTensor, torch.FloatTensor]#

Conditions on the size of the states/actions and calls the appropriate singular or parallel dynamics.

Parameters:

  • states (torch.FloatTensor) – (nx,) or (num_envs, nx) sized vector of states

  • actions (torch.FloatTensor) – (2, na) or (num_envs, 2, na) array containing the start and end action vectors of the desired trajectory.

  • action_time_step (float) – the hold time for the action.

Returns:

A tuple of (next state, intermediate states)

Return type:

Tuple[torch.FloatTensor, torch.FloatTensor]

get_num_substeps(time_step: float | None = None) → int#
Parameters:

time_step (Optional[float]) –

Return type:

int

get_gradient_placeholders(size: int | None = None) → Tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike]#
Parameters:

size (Optional[int]) –

Return type:

Tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike, numpy.typing.ArrayLike]

get_sub_stepped_gradients(state: torch.FloatTensor, action: torch.FloatTensor, num_substeps: int | None = None) → Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]#
Parameters:

  • state (torch.FloatTensor) –

  • action (torch.FloatTensor) –

  • num_substeps (Optional[int]) –

Return type:

Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]

get_gradients(states: torch.FloatTensor, actions: torch.FloatTensor) → Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]#

Computes the dynamics gradients.

Parameters:

  • states (torch.FloatTensor) – (nx,) or (num_envs, nx) sized vector of states

  • actions (torch.FloatTensor) – (na,) or (num_envs, na) array containing the action vectors.

Returns:

state_gradients_state – (nx, nx) or (num_envs, nx, nx), state_gradients_control: (nx, nu) or (num_envs, nx, nu), sensor_gradients_state: (ns, nx) or (num_envs, ns, nx), sensor_gradients_control: (ns, nu) or (num_envs, ns, nu),

Return type:

Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]

set_state(state: torch.FloatTensor) → None#
Parameters:

state (torch.FloatTensor) –

Return type:

None

get_state() → torch.FloatTensor#
Return type:

torch.FloatTensor

set_action(action: torch.FloatTensor) → None#
Parameters:

action (torch.FloatTensor) –

Return type:

None

get_action() → torch.FloatTensor#
Return type:

torch.FloatTensor

update_sensor() → None#
Return type:

None

get_sensor(states: torch.FloatTensor) → torch.FloatTensor#

Update the sensor measurement held by plant.data. This only supports position- and velocity-based sensors, NOT ACCLERATION-BASED sensors. We use the minimal set of computations extracted from mj_step1, see the link below for more details: https://mujoco.readthedocs.io/en/latest/programming/simulation.html?highlight=mj_step1#simulation-loop Finally, returns the sensor measurement.

Parameters:

states (torch.FloatTensor) – (nx,) or (num_envs, nx) sized vector of states

Returns:

sensor (nsensordata,) or (num_envs, nsensordata)

Return type:

torch.FloatTensor

state_difference(s1: torch.FloatTensor, s2: torch.FloatTensor, h: float = 1.0) → torch.FloatTensor#

Compute finite-difference velocity given two state vectors and a time step h ds = (s2 - s1) / h

Parameters:

  • s1 (torch.FloatTensor) –

  • s2 (torch.FloatTensor) –

  • h (float) –

Return type:

torch.FloatTensor

state_addition(s1: torch.FloatTensor, ds: torch.FloatTensor, h: float = 1.0) → torch.FloatTensor#

Integrate forward a state s with a velocity ds for a time step h. s2 = s1 + h * ds

Parameters:

  • s1 (torch.FloatTensor) –

  • ds (torch.FloatTensor) –

  • h (float) –

Return type:

torch.FloatTensor

get_mass() → float#
Return type:

float

get_quat_indices() → None#

Stores the indices of the state corresponding to quaternion in class field.

Return type:

None

normalize_state(states: torch.FloatTensor) → torch.FloatTensor#
Parameters:

states (torch.FloatTensor) – (num_envs, nx) tensor of randomly sampled states

Return type:

torch.FloatTensor

Outputs:

torch.FloatTensor: (num_envs, nx) tensor with quaternion portions normalized

get_collision_free(states: torch.FloatTensor) → torch.FloatTensor | None#
Parameters:

  • state (FloatTensor) – (num_envs, nx) tensor of randomly sampled states

  • states (torch.FloatTensor) –

Return type:

Optional[torch.FloatTensor]

Outputs:

Optional[FloatTensor]: tensor of all collision free states (or None if none exist)