Functions — github.com/LeCAR-Lab/model-based-diffusion

CodeHub MCPcopy Create free account

hub / github.com/LeCAR-Lab/model-based-diffusion / functions

Functions93 in github.com/LeCAR-Lab/model-based-diffusion

⨍ Functions 93 ◇ Types & classes 14

↓ 4 callersFunctionobjective_function

Define the objective function J(x) with multiple local minima. Args: x: Input values Returns: Objective function values

mbd/notebooks/01_1d_demo.py:27

↓ 4 callersMethodrender

(self, ax, xs: jnp.ndarray)

mbd/envs/car2d.py:112

↓ 3 callersMethod_get_reward

(self, pipeline_state: pipeline.State)

mbd/envs/pushT.py:51

↓ 2 callersMethod_get_done

(self, pipeline_state: pipeline.State)

mbd/envs/pushT.py:64

↓ 2 callersMethod_get_obs

(self, pipeline_state: pipeline.State)

mbd/envs/pushT.py:48

↓ 2 callersMethod_get_obs

(self, pipeline_state: base.State, action: jax.Array)

mbd/envs/humanoidstandup.py:47

↓ 2 callersMethod_get_obs

Observe cartpole body position and velocities.

mbd/envs/cartpole.py:54

↓ 2 callersMethod_get_obs

Returns the environment observations.

mbd/envs/walker2d.py:50

↓ 2 callersMethod_get_obs

Returns the environment observations.

mbd/envs/hopper.py:49

↓ 2 callersMethod_get_obs

(self, pipeline_state: base.State)

mbd/envs/humanoidtrack.py:84

↓ 2 callersMethod_get_obs

(self, pipeline_state: base.State, action: jax.Array)

mbd/envs/humanoidrun.py:43

↓ 2 callersFunction_one_hot

Create a one-hot encoding of x of size k.

mbd/blackbox/mbd_mnist.py:42

↓ 2 callersFunction_partial_flatten

Flatten all but the first dimension of an ndarray.

mbd/blackbox/mbd_mnist.py:37

↓ 2 callersFunctionaccuracy

(params, batch)

mbd/blackbox/mbd_mnist.py:164

↓ 2 callersFunctionparse_images

mbd/blackbox/mbd_mnist.py:57

↓ 2 callersFunctionparse_labels

mbd/blackbox/mbd_mnist.py:52

↓ 2 callersFunctionrender_us

mbd/scripts/vis_diffusion.py:115

↓ 2 callersFunctionreverse_once

(carry, unused)

mbd/blackbox/mbd_opt.py:65

↓ 2 callersFunctionreverse_once

(carry, unused)

mbd/blackbox/mbd_mnist.py:194

↓ 1 callersFunction_download

Download a url to a file in the JAX data temp directory.

mbd/blackbox/mbd_mnist.py:27

↓ 1 callersMethod_get_reward

(self, pipeline_state: base.State)

mbd/envs/humanoidstandup.py:50

↓ 1 callersMethod_get_reward

Returns the environment reward.

mbd/envs/walker2d.py:58

↓ 1 callersMethod_get_reward

Returns the environment reward.

mbd/envs/hopper.py:57

↓ 1 callersMethod_get_reward

mbd/envs/humanoidtrack.py:87

↓ 1 callersMethod_get_reward

(self, pipeline_state: base.State)

mbd/envs/humanoidrun.py:46

↓ 1 callersFunctionadd_noise_batch_to_params

(params, sigma, rng)

mbd/blackbox/mbd_mnist.py:123

↓ 1 callersFunctioncheck_collision

(x, obs_center, obs_radius)

mbd/envs/car2d.py:30

↓ 1 callersFunctiondumps

mbd/scripts/vis_diffusion.py:27

↓ 1 callersFunctioneval_us

(step_env, state, us)

↓ 1 callersFunctionget_params_batch_weighted_sum

(params_batch, weights)

mbd/blackbox/mbd_mnist.py:145

↓ 1 callersMethodget_reward

mbd/envs/car2d.py:89

↓ 1 callersFunctionloss

(params, batch)

mbd/blackbox/mbd_mnist.py:158

↓ 1 callersFunctionmain

mbd/envs/pushT.py:77

↓ 1 callersFunctionmnist

Download, parse and process MNIST data to unit scale and one-hot labels.

mbd/blackbox/mbd_mnist.py:80

↓ 1 callersFunctionmnist_raw

Download and parse the raw MNIST dataset.

mbd/blackbox/mbd_mnist.py:47

↓ 1 callersMethodreset

(self, rng: jnp.ndarray)

mbd/envs/pushT.py:22

↓ 1 callersFunctionreverse

mbd/planners/mbd_planner.py:138

↓ 1 callersFunctionreverse_once

(carry, unused)

mbd/planners/mbd_planner.py:98

↓ 1 callersFunctionrk4

(dynamics, x, u, dt)

mbd/envs/car2d.py:22

↓ 1 callersFunctionrollout_us

(step_env, state, us)

mbd/utils.py:14

↓ 1 callersFunctionrun_diffusion

mbd/planners/mbd_planner.py:38

↓ 1 callersFunctionrun_exp

mbd/blackbox/mbd_opt.py:83

↓ 1 callersFunctionrun_multiple_seed

mbd/scripts/run_mbd.py:17

↓ 1 callersFunctionrun_multiple_temp

mbd/scripts/run_mbd.py:42

↓ 1 callersFunctionrun_path_integral

mbd/planners/path_integral.py:55

↓ 1 callersFunctionupdate

mbd/planners/path_integral.py:130

↓ 1 callersFunctionupdate_once

(carry, unused)

mbd/planners/path_integral.py:112

(self, backend: str = "generalized")

mbd/envs/pushT.py:16

mbd/envs/humanoidstandup.py:14

(self, backend="positional", **kwargs)

mbd/envs/cartpole.py:12

mbd/envs/walker2d.py:13

mbd/envs/hopper.py:12

(self, mode="jog")

mbd/envs/humanoidtrack.py:15

mbd/envs/humanoidrun.py:14

mbd/envs/car2d.py:44

mbd/blackbox/mbd_opt.py:32

Methodaction_size

mbd/envs/pushT.py:69

Methodaction_size

mbd/envs/cartpole.py:51

Methodaction_size

mbd/envs/car2d.py:105

Functionadd_noise_to_params

(params, sigma, rng)

mbd/blackbox/mbd_mnist.py:175

Functioncar_dynamics

mbd/envs/car2d.py:10

Functioncem_update

(weights, Y0s, sigma, mu_0t)

mbd/planners/path_integral.py:49

Functioncma_es_update

(weights, Y0s, sigma, mu_0t)

mbd/planners/path_integral.py:40

Methodconstruct

mbd/scripts/vis_manim.py:7

Methoddemo_vector_field_fn

mbd/scripts/vis_manim.py:44

Functioneval_fn

mbd/blackbox/mbd_mnist.py:171

Methodeval_xref_logpd

(self, pipeline_state)

mbd/envs/humanoidtrack.py:99

Methodeval_xref_logpd

mbd/envs/car2d.py:96

Functionget_env

(env_name: str)

mbd/envs/__init__.py:13

mbd/blackbox/mbd_opt.py:44

Methodobservation_size

mbd/envs/pushT.py:73

Methodobservation_size

mbd/envs/car2d.py:109

Functionprogress

(num_steps, metrics)

mbd/rl/train_brax.py:182

Functionrastrigin

mbd/blackbox/mbd_opt.py:39

Functionrender_us

(step_env, sys, state, us)

mbd/utils.py:23

Resets the environment to an initial state.

mbd/envs/humanoidstandup.py:19

Resets the environment to an initial state.

mbd/envs/cartpole.py:25

Resets the environment to an initial state.

mbd/envs/walker2d.py:21

Resets the environment to an initial state.

mbd/envs/hopper.py:20

Resets the environment to an initial state.

mbd/envs/humanoidtrack.py:48

Resets the environment to an initial state.

mbd/envs/humanoidrun.py:19

Resets the environment to an initial state.

mbd/envs/car2d.py:73

Functionsoftmax_update

(weights, Y0s, sigma, mu_0t)

mbd/planners/path_integral.py:34

Methodsquare_vector_field_fn

mbd/scripts/vis_manim.py:296

(self, state: State, action: jnp.ndarray)

mbd/envs/pushT.py:39

Runs one timestep of the environment's dynamics.

mbd/envs/humanoidstandup.py:38

Run one timestep of the environment's dynamics.

mbd/envs/cartpole.py:40

Runs one timestep of the environment's dynamics.

mbd/envs/walker2d.py:37

Runs one timestep of the environment's dynamics.

mbd/envs/hopper.py:36

Runs one timestep of the environment's dynamics.

mbd/envs/humanoidtrack.py:63

Runs one timestep of the environment's dynamics.

mbd/envs/humanoidrun.py:34

Run one timestep of the environment's dynamics.

mbd/envs/car2d.py:78