Functions — github.com/NVlabs/neural-robot-dynamics

Method__init__

(self, name)

utils/time_report.py:21

Method__init__

(self, cuda_synchronize = False)

utils/time_report.py:50

Method__init__

( self, # warp environment arguments env_name, num_envs, warp_env_cfg

envs/neural_environment.py:45

Method__init__

(self, env: Environment)

envs/abstract_contact_environment.py:208

Method__init__

( self, env: NeuralEnvironment, render_mode: str = None, max_episode_length: i

envs/rlgames_env_wrapper.py:61

Method__init__

(self)

envs/rlgames_env_wrapper.py:312

Method__init__

(self, num_contacts_per_env, num_envs, model: wp.sim.Model, device = 'cuda:0')

envs/abstract_contact.py:40

Method__init__

( self, seed=42, random_reset=True, task="run", obs_type="dflex",

envs/warp_sim_envs/env_ant.py:343

Method__init__

( self, num_envs: int = None, episode_frames: int = None, integrator_type: Int

envs/warp_sim_envs/environment.py:260

Method__init__

(self, seed = 42, random_reset = True, **kwargs)

envs/warp_sim_envs/env_cartpole.py:164

Method__init__

(self, **kwargs)

envs/warp_sim_envs/env_anymal_joint_position_control.py:60

Method__init__

(self, seed=42, random_reset=True, **kwargs)

envs/warp_sim_envs/env_pendulum_with_contact.py:88

Method__init__

( self, seed=42, control_mode="joint_position_control", random_reset=False,

envs/warp_sim_envs/env_franka_panda.py:250

Method__init__

( self, seed=42, random_reset=True, task="forward", obs_type="dflex",

envs/warp_sim_envs/env_anymal.py:429

Method__init__

(self, seed=42, random_reset=True, **kwargs)

envs/warp_sim_envs/env_cube_toss.py:191

Method__init__

( self, num_states_history = 1, **kwargs )

integrators/integrator_neural_stateful.py:30

Method__init__

( self, **kwargs )

integrators/integrator_neural_transformer.py:29

Method__init__

( self, reset_seq_length = 1, **kwargs )

integrators/integrator_neural_rnn.py:26

Method__init__

Args: model (Model): The Warp sim model to use for simulation which stores static information about the system to

integrators/integrator_neural.py:134

Method__init__

(self, ndim, bias)

models/model_transformer.py:22

Method__init__

(self, config)

models/model_transformer.py:32

Method__init__

(self, config)

models/model_transformer.py:106

Method__init__

(self, config)

models/model_transformer.py:130

Method__init__

( self, input_shape, output_dim, network_cfg, device='cuda:0' )

models/models.py:23

Method__init__

(self, in_features, network_cfg, device='cuda:0')

models/base_models.py:50

Method__init__

(self, in_features, network_cfg, device='cuda:0')

models/base_models.py:126

Method__init__

( self, neural_env, cfg, model_checkpoint_path=None, device='cuda:0'

algorithms/sequence_model_trainer.py:24

Method__init__

( self, neural_env: NeuralEnvironment, cfg: dict, model_checkpoint_path: Op

algorithms/vanilla_trainer.py:47

Method__init__

( self, seed=None, scramble=False, )

generate/simulation_sampler.py:83

Method__init__

(self)

generate/simulation_sampler.py:107

Method__init__

( self, env: NeuralEnvironment, joint_q_min: Union[float, np.ndarray], joint_q

generate/simulation_sampler.py:124

Method__init__

( self, env: NeuralEnvironment, joint_q_min: Union[float, np.ndarray], joint_q

generate/trajectory_sampler.py:39

Method__init__

( self, env: NeuralEnvironment, joint_q_min: Union[float, np.ndarray], joint_q

generate/trajectory_sampler_pendulum.py:44

Method__len__

(self)

utils/datasets.py:92

Method__len__

(self)

utils/datasets.py:174

Method__str__

(self)

envs/warp_sim_envs/environment.py:45

Method__str__

(self)

envs/warp_sim_envs/environment.py:55

Function_acquire_states

( joint_q: wp.array(dtype=float), joint_qd: wp.array(dtype=float), q_count: int, qd_count: int

utils/warp_utils.py:34

Function_assign_joint_acts_from_actions

( actions: wp.array(dtype=float, ndim=2), control_count: int, joint_act_dim: int, control_dofs

utils/warp_utils.py:49

Function_assign_states

( states: wp.array(dtype=float, ndim=2), q_count: int, qd_count: int, # output joint_q: wp

utils/warp_utils.py:19

Method_init_weights

(self, module)

models/model_transformer.py:175

Method_update_states

(self, model: Model, warp_states: State, joint_act)

integrators/integrator_neural_stateful.py:74

Methodabstract_contacts

(self)

envs/neural_environment.py:225

Functionacquire_states_to_torch

(warp_env, torch_states)

utils/warp_utils.py:84

Methodaction_dim

(self)

envs/neural_environment.py:192

Methodaction_dim

(self)

generate/simulation_sampler.py:232

Methodaction_limits

(self)

envs/neural_environment.py:196

Methodafter_clear_stats

(self)

envs/rlgames_env_wrapper.py:343

Methodafter_init

(self, algo)

envs/rlgames_env_wrapper.py:315

Methodafter_print_stats

(self, frame, epoch_num, total_time)

envs/rlgames_env_wrapper.py:346

Methodafter_reset

(self)

envs/warp_sim_envs/environment.py:880

Methodafter_simulate

(self)

envs/warp_sim_envs/environment.py:498

Functionangle_normalize

(x: float)

envs/warp_sim_envs/utils.py:30

Functionant_running_cost

( joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), inv_start_rot: wp.qua

envs/warp_sim_envs/env_ant.py:30

Functionant_spinning_cost

( joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), inv_start_rot: wp.qua

envs/warp_sim_envs/env_ant.py:77

Functionant_spinning_tracking_cost

( joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), inv_start_rot: wp.qua

envs/warp_sim_envs/env_ant.py:116

Functionanymal_forward_cost

( joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), contact_depths: wp.ar

envs/warp_sim_envs/env_anymal.py:28

Functionanymal_side_cost

( joint_q: wp.array(dtype=wp.float32), joint_qd: wp.array(dtype=wp.float32), contact_depths: wp.ar

envs/warp_sim_envs/env_anymal.py:98

Functionapply_joint_position_pd_control

( actions: wp.array(dtype=wp.float32, ndim=1), action_scale: wp.float32, baseline_joint_q: wp.arra

envs/warp_sim_envs/env_anymal_joint_position_control.py:22

Functionapply_joint_position_pd_control

( actions: wp.array(dtype=wp.float32, ndim=1), action_scale: wp.float32, baseline_joint_q: wp.arra

envs/warp_sim_envs/env_franka_panda.py:27

Methodassign_control

( self, actions: wp.array, control: wp.sim.Control, state: wp.sim.State )

envs/warp_sim_envs/env_anymal_joint_position_control.py:75

Methodassign_control

( self, actions: wp.array, control: wp.sim.Control, state: wp.sim.State, )

envs/warp_sim_envs/env_franka_panda.py:525

Methodassign_control

( self, actions: wp.array, control: wp.sim.Control, state: wp.sim.State, )

envs/warp_sim_envs/env_anymal.py:676

Functionassign_controls

( actions: wp.array(dtype=wp.float32, ndim=2), gains: wp.array(dtype=wp.float32), limits: wp.array

envs/warp_sim_envs/utils.py:352

Functionassign_joint_acts_from_actions_torch

(warp_env, actions)

utils/warp_utils.py:99

Functionassign_joint_q_qd_obs

( joint_q: wp.array(dtype=float), joint_qd: wp.array(dtype=float), dof_q_per_env: int, dof_qd_

envs/warp_sim_envs/environment.py:103

Functionassign_states_from_torch

(warp_env, torch_states)

utils/warp_utils.py:67

Methodbefore_model_forward

(self)

integrators/integrator_neural_rnn.py:39

Methodbefore_simulate

(self)

envs/warp_sim_envs/environment.py:495

Methodbodies_per_env

(self)

envs/neural_environment.py:176

Functionclear_nonfinite_values

(values: wp.array(dtype=float, ndim=2))

envs/warp_sim_envs/wrapper_utils.py:95

Methodclose

Clean up the environment's resources.

envs/rlgames_env_wrapper.py:232

Methodclose

(self)

envs/warp_sim_envs/environment.py:883

Functioncollate_fn_BatchTransitionDataset

(batch)

utils/datasets.py:111

Functioncollision_detection_ground

( body_q: wp.array(dtype=wp.transform), shape_X_bs: wp.array(dtype=wp.transform), shape_body: wp.a

envs/abstract_contact_environment.py:134

Functioncompute_body_jacobian

( model: wp.sim.Model, joint_q: wp.array, joint_qd: wp.array, body_id: Union[int, str], of

envs/warp_sim_envs/utils.py:693

Functioncompute_body_out

( body_qd: wp.array(dtype=wp.spatial_vector), body_out: wp.array(dtype=float) )

envs/warp_sim_envs/utils.py:713

Functioncompute_contact_points_0_world

( body_q: wp.array(dtype=wp.transform), shape_body: wp.array(dtype=int), contact_shape0: wp.array(

generate/simulation_sampler.py:32

Functioncompute_contact_points_1

( body_q: wp.array(dtype=wp.transform), shape_body: wp.array(dtype=int), contact_shape0: wp.array(

generate/simulation_sampler.py:50

Methodcompute_cost_termination

( self, state: wp.sim.State, control: wp.sim.Control, step: int, traj_

envs/warp_sim_envs/env_ant.py:466

Methodcompute_cost_termination

( self, state: wp.sim.State, control: wp.sim.Control, step: int, max_e

envs/warp_sim_envs/environment.py:649

Methodcompute_cost_termination

( self, state: wp.sim.State, control: wp.sim.Control, step: int, traj_

envs/warp_sim_envs/env_cartpole.py:199

Methodcompute_cost_termination

( self, state: wp.sim.State, control: wp.sim.Control, step: int, max_e

envs/warp_sim_envs/env_franka_panda.py:457

Methodcompute_cost_termination

( self, state: wp.sim.State, control: wp.sim.Control, step: int, traj_

envs/warp_sim_envs/env_anymal.py:577

Methodcompute_observations