SYMBOL INDEX (590 symbols across 56 files)

FILE: data_gen_scripts/generate_antsoccer.py
  function load_agent (line 31) | def load_agent(restore_path, restore_epoch, ob_dim, action_dim):
  function main (line 52) | def main(_):

FILE: data_gen_scripts/generate_locomaze.py
  function main (line 29) | def main(_):

FILE: data_gen_scripts/generate_manipspace.py
  function main (line 33) | def main(_):

FILE: data_gen_scripts/generate_powderworld.py
  function main (line 23) | def main(_):

FILE: data_gen_scripts/main_sac.py
  function main (line 49) | def main(_):

FILE: data_gen_scripts/online_env_utils.py
  function make_online_env (line 5) | def make_online_env(env_name):

FILE: data_gen_scripts/viz_utils.py
  function get_2d_colors (line 7) | def get_2d_colors(points, min_point, max_point):
  function visualize_trajs (line 21) | def visualize_trajs(env_name, trajs):

FILE: impls/agents/crl.py
  class CRLAgent (line 13) | class CRLAgent(flax.struct.PyTreeNode):
    method contrastive_loss (line 24) | def contrastive_loss(self, batch, grad_params, module_name='critic'):
    method actor_loss (line 71) | def actor_loss(self, batch, grad_params, rng=None):
    method total_loss (line 136) | def total_loss(self, batch, grad_params, rng=None):
    method update (line 161) | def update(self, batch):
    method sample_actions (line 173) | def sample_actions(
    method create (line 188) | def create(
  function get_config (line 292) | def get_config():

FILE: impls/agents/gcbc.py
  class GCBCAgent (line 13) | class GCBCAgent(flax.struct.PyTreeNode):
    method actor_loss (line 20) | def actor_loss(self, batch, grad_params, rng=None):
    method total_loss (line 42) | def total_loss(self, batch, grad_params, rng=None):
    method update (line 56) | def update(self, batch):
    method sample_actions (line 68) | def sample_actions(
    method create (line 83) | def create(
  function get_config (line 143) | def get_config():

FILE: impls/agents/gciql.py
  class GCIQLAgent (line 14) | class GCIQLAgent(flax.struct.PyTreeNode):
    method expectile_loss (line 25) | def expectile_loss(adv, diff, expectile):
    method value_loss (line 30) | def value_loss(self, batch, grad_params):
    method critic_loss (line 44) | def critic_loss(self, batch, grad_params):
    method actor_loss (line 61) | def actor_loss(self, batch, grad_params, rng=None):
    method total_loss (line 126) | def total_loss(self, batch, grad_params, rng=None):
    method target_update (line 147) | def target_update(self, network, module_name):
    method update (line 157) | def update(self, batch):
    method sample_actions (line 170) | def sample_actions(
    method create (line 185) | def create(
  function get_config (line 276) | def get_config():

FILE: impls/agents/gcivl.py
  class GCIVLAgent (line 14) | class GCIVLAgent(flax.struct.PyTreeNode):
    method expectile_loss (line 25) | def expectile_loss(adv, diff, expectile):
    method value_loss (line 30) | def value_loss(self, batch, grad_params):
    method actor_loss (line 63) | def actor_loss(self, batch, grad_params, rng=None):
    method total_loss (line 95) | def total_loss(self, batch, grad_params, rng=None):
    method target_update (line 112) | def target_update(self, network, module_name):
    method update (line 122) | def update(self, batch):
    method sample_actions (line 135) | def sample_actions(
    method create (line 150) | def create(
  function get_config (line 223) | def get_config():

FILE: impls/agents/hiql.py
  class HIQLAgent (line 14) | class HIQLAgent(flax.struct.PyTreeNode):
    method expectile_loss (line 22) | def expectile_loss(adv, diff, expectile):
    method value_loss (line 27) | def value_loss(self, batch, grad_params):
    method low_actor_loss (line 60) | def low_actor_loss(self, batch, grad_params):
    method high_actor_loss (line 99) | def high_actor_loss(self, batch, grad_params):
    method total_loss (line 127) | def total_loss(self, batch, grad_params, rng=None):
    method target_update (line 146) | def target_update(self, network, module_name):
    method update (line 156) | def update(self, batch):
    method sample_actions (line 169) | def sample_actions(
    method create (line 195) | def create(
  function get_config (line 319) | def get_config():

FILE: impls/agents/qrl.py
  class QRLAgent (line 14) | class QRLAgent(flax.struct.PyTreeNode):
    method value_loss (line 30) | def value_loss(self, batch, grad_params):
    method dynamics_loss (line 60) | def dynamics_loss(self, batch, grad_params):
    method actor_loss (line 78) | def actor_loss(self, batch, grad_params, rng=None):
    method total_loss (line 146) | def total_loss(self, batch, grad_params, rng=None):
    method update (line 171) | def update(self, batch):
    method sample_actions (line 183) | def sample_actions(
    method create (line 198) | def create(
  function get_config (line 294) | def get_config():

FILE: impls/agents/sac.py
  class SACAgent (line 13) | class SACAgent(flax.struct.PyTreeNode):
    method critic_loss (line 20) | def critic_loss(self, batch, grad_params, rng):
    method actor_loss (line 44) | def actor_loss(self, batch, grad_params, rng):
    method total_loss (line 80) | def total_loss(self, batch, grad_params, rng=None):
    method target_update (line 98) | def target_update(self, network, module_name):
    method update (line 108) | def update(self, batch):
    method sample_actions (line 121) | def sample_actions(
    method create (line 135) | def create(
  function get_config (line 198) | def get_config():

FILE: impls/main.py
  function main (line 45) | def main(_):

FILE: impls/utils/datasets.py
  function get_size (line 11) | def get_size(data):
  function random_crop (line 18) | def random_crop(img, crop_from, padding):
  function batched_random_crop (line 31) | def batched_random_crop(imgs, crop_froms, padding):
  class Dataset (line 36) | class Dataset(FrozenDict):
    method create (line 46) | def create(cls, freeze=True, **fields):
    method __init__ (line 59) | def __init__(self, *args, **kwargs):
    method get_random_idxs (line 65) | def get_random_idxs(self, num_idxs):
    method sample (line 72) | def sample(self, batch_size, idxs=None):
    method get_subset (line 78) | def get_subset(self, idxs):
  class ReplayBuffer (line 86) | class ReplayBuffer(Dataset):
    method create (line 93) | def create(cls, transition, size):
    method create_from_initial_dataset (line 109) | def create_from_initial_dataset(cls, init_dataset, size):
    method __init__ (line 127) | def __init__(self, *args, **kwargs):
    method add_transition (line 134) | def add_transition(self, transition):
    method clear (line 144) | def clear(self):
  class GCDataset (line 150) | class GCDataset:
    method __post_init__ (line 182) | def __post_init__(self):
    method sample (line 205) | def sample(self, batch_size, idxs=None, evaluation=False):
    method sample_goals (line 252) | def sample_goals(self, idxs, p_curgoal, p_trajgoal, p_randomgoal, geom...
    method augment (line 283) | def augment(self, batch, keys):
    method get_observations (line 295) | def get_observations(self, idxs):
    method get_stacked_observations (line 302) | def get_stacked_observations(self, idxs):
  class HGCDataset (line 313) | class HGCDataset(GCDataset):
    method sample (line 321) | def sample(self, batch_size, idxs=None, evaluation=False):

FILE: impls/utils/encoders.py
  class ResnetStack (line 10) | class ResnetStack(nn.Module):
    method __call__ (line 18) | def __call__(self, x):
  class ImpalaEncoder (line 60) | class ImpalaEncoder(nn.Module):
    method setup (line 70) | def setup(self):
    method __call__ (line 83) | def __call__(self, x, train=True, cond_var=None):
  class GCEncoder (line 103) | class GCEncoder(nn.Module):
    method __call__ (line 116) | def __call__(self, observations, goals=None, goal_encoded=False):

FILE: impls/utils/env_utils.py
  class EpisodeMonitor (line 14) | class EpisodeMonitor(gymnasium.Wrapper):
    method __init__ (line 17) | def __init__(self, env):
    method _reset_stats (line 22) | def _reset_stats(self):
    method step (line 27) | def step(self, action):
    method reset (line 43) | def reset(self, *args, **kwargs):
  class FrameStackWrapper (line 48) | class FrameStackWrapper(gymnasium.Wrapper):
    method __init__ (line 51) | def __init__(self, env, num_stack):
    method get_observation (line 61) | def get_observation(self):
    method reset (line 65) | def reset(self, **kwargs):
    method step (line 73) | def step(self, action):
  function make_env_and_datasets (line 79) | def make_env_and_datasets(dataset_name, frame_stack=None):

FILE: impls/utils/evaluation.py
  function supply_rng (line 8) | def supply_rng(f, rng=jax.random.PRNGKey(0)):
  function flatten (line 19) | def flatten(d, parent_key='', sep='.'):
  function add_to (line 31) | def add_to(dict_of_lists, single_dict):
  function evaluate (line 37) | def evaluate(

FILE: impls/utils/flax_utils.py
  class ModuleDict (line 16) | class ModuleDict(nn.Module):
    method __call__ (line 28) | def __call__(self, *args, name=None, **kwargs):
  class TrainState (line 53) | class TrainState(flax.struct.PyTreeNode):
    method create (line 73) | def create(cls, model_def, params, tx=None, **kwargs):
    method __call__ (line 90) | def __call__(self, *args, params=None, method=None, **kwargs):
    method select (line 116) | def select(self, name):
    method apply_gradients (line 120) | def apply_gradients(self, grads, **kwargs):
    method apply_loss_fn (line 132) | def apply_loss_fn(self, loss_fn):
  function save_agent (line 162) | def save_agent(agent, save_dir, epoch):
  function restore_agent (line 181) | def restore_agent(agent, restore_path, restore_epoch):

FILE: impls/utils/log_utils.py
  class CsvLogger (line 12) | class CsvLogger:
    method __init__ (line 15) | def __init__(self, path):
    method log (line 21) | def log(self, row, step):
    method close (line 35) | def close(self):
  function get_exp_name (line 40) | def get_exp_name(seed):
  function get_flag_dict (line 53) | def get_flag_dict():
  function setup_wandb (line 62) | def setup_wandb(
  function reshape_video (line 94) | def reshape_video(v, n_cols=None):
  function get_wandb_video (line 116) | def get_wandb_video(renders=None, n_cols=None, fps=15):

FILE: impls/utils/networks.py
  function default_init (line 10) | def default_init(scale=1.0):
  function ensemblize (line 15) | def ensemblize(cls, num_qs, out_axes=0, **kwargs):
  class Identity (line 28) | class Identity(nn.Module):
    method __call__ (line 31) | def __call__(self, x):
  class MLP (line 35) | class MLP(nn.Module):
    method __call__ (line 53) | def __call__(self, x):
  class LengthNormalize (line 63) | class LengthNormalize(nn.Module):
    method __call__ (line 70) | def __call__(self, x):
  class Param (line 74) | class Param(nn.Module):
    method __call__ (line 80) | def __call__(self):
  class LogParam (line 84) | class LogParam(nn.Module):
    method __call__ (line 90) | def __call__(self):
  class TransformedWithMode (line 95) | class TransformedWithMode(distrax.Transformed):
    method mode (line 98) | def mode(self):
  class RunningMeanStd (line 102) | class RunningMeanStd(flax.struct.PyTreeNode):
    method normalize (line 119) | def normalize(self, batch):
    method unnormalize (line 124) | def unnormalize(self, batch):
    method update (line 127) | def update(self, batch):
  class GCActor (line 143) | class GCActor(nn.Module):
    method setup (line 168) | def setup(self):
    method __call__ (line 177) | def __call__(
  class GCDiscreteActor (line 219) | class GCDiscreteActor(nn.Module):
    method setup (line 234) | def setup(self):
    method __call__ (line 238) | def __call__(
  class GCValue (line 269) | class GCValue(nn.Module):
    method setup (line 286) | def setup(self):
    method __call__ (line 294) | def __call__(self, observations, goals=None, actions=None):
  class GCDiscreteCritic (line 317) | class GCDiscreteCritic(GCValue):
    method __call__ (line 322) | def __call__(self, observations, goals=None, actions=None):
  class GCBilinearValue (line 327) | class GCBilinearValue(nn.Module):
    method setup (line 351) | def setup(self):
    method __call__ (line 359) | def __call__(self, observations, goals, actions=None, info=False):
  class GCDiscreteBilinearCritic (line 392) | class GCDiscreteBilinearCritic(GCBilinearValue):
    method __call__ (line 397) | def __call__(self, observations, goals=None, actions=None, info=False):
  class GCMRNValue (line 402) | class GCMRNValue(nn.Module):
    method setup (line 420) | def setup(self):
    method __call__ (line 423) | def __call__(self, observations, goals, is_phi=False, info=False):
  class GCIQEValue (line 456) | class GCIQEValue(nn.Module):
    method setup (line 475) | def setup(self):
    method __call__ (line 479) | def __call__(self, observations, goals, is_phi=False, info=False):

FILE: ogbench/locomaze/ant.py
  class AntEnv (line 10) | class AntEnv(MujocoEnv, utils.EzPickle):
    method __init__ (line 26) | def __init__(
    method step (line 69) | def step(self, action):
    method get_ob (line 97) | def get_ob(self):
    method reset_model (line 103) | def reset_model(self):
    method get_xy (line 115) | def get_xy(self):
    method set_xy (line 118) | def set_xy(self, xy):

FILE: ogbench/locomaze/humanoid.py
  class HumanoidEnv (line 12) | class HumanoidEnv(MujocoEnv, utils.EzPickle):
    method __init__ (line 27) | def __init__(
    method step (line 65) | def step(self, action):
    method _step_mujoco_simulation (line 93) | def _step_mujoco_simulation(self, ctrl, n_frames):
    method get_ob (line 106) | def get_ob(self):
    method disable (line 134) | def disable(self, *flags):
    method reset_model (line 150) | def reset_model(self):
    method get_xy (line 170) | def get_xy(self):
    method set_xy (line 173) | def set_xy(self, xy):

FILE: ogbench/locomaze/maze.py
  function make_maze_env (line 13) | def make_maze_env(loco_env_type, maze_env_type, *args, **kwargs):

FILE: ogbench/locomaze/point.py
  class PointEnv (line 11) | class PointEnv(MujocoEnv, utils.EzPickle):
    method __init__ (line 26) | def __init__(
    method step (line 64) | def step(self, action):
    method get_ob (line 97) | def get_ob(self):
    method reset_model (line 100) | def reset_model(self):
    method get_xy (line 108) | def get_xy(self):
    method set_xy (line 111) | def set_xy(self, xy):

FILE: ogbench/manipspace/controllers/diff_ik.py
  function angle_diff (line 8) | def angle_diff(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
  class DiffIKController (line 12) | class DiffIKController:
    method __init__ (line 15) | def __init__(
    method _forward_kinematics (line 41) | def _forward_kinematics(self) -> None:
    method _integrate (line 46) | def _integrate(self, update: np.ndarray) -> None:
    method _compute_translational_error (line 50) | def _compute_translational_error(self, pos: np.ndarray) -> None:
    method _compute_rotational_error (line 54) | def _compute_rotational_error(self, quat: np.ndarray) -> None:
    method _compute_jacobian (line 62) | def _compute_jacobian(self) -> None:
    method _error_threshold_reached (line 69) | def _error_threshold_reached(self, pos_thresh: float, ori_thresh: floa...
    method _solve (line 75) | def _solve(self) -> np.ndarray:
    method _scale_update (line 85) | def _scale_update(self, update: np.ndarray) -> np.ndarray:
    method solve (line 92) | def solve(

FILE: ogbench/manipspace/envs/cube_env.py
  class CubeEnv (line 9) | class CubeEnv(ManipSpaceEnv):
    method __init__ (line 17) | def __init__(self, env_type, permute_blocks=True, *args, **kwargs):
    method set_tasks (line 75) | def set_tasks(self):
    method add_objects (line 510) | def add_objects(self, arena_mjcf):
    method post_compilation_objects (line 550) | def post_compilation_objects(self):
    method initialize_episode (line 563) | def initialize_episode(self):
    method set_new_target (line 647) | def set_new_target(self, return_info=True, p_stack=0.5):
    method _compute_successes (line 709) | def _compute_successes(self):
    method post_step (line 722) | def post_step(self):
    method add_object_info (line 746) | def add_object_info(self, ob_info):
    method compute_observation (line 766) | def compute_observation(self):
    method compute_oracle_observation (line 796) | def compute_oracle_observation(self):
    method compute_reward (line 808) | def compute_reward(self):

FILE: ogbench/manipspace/envs/env.py
  class CustomMuJoCoEnv (line 14) | class CustomMuJoCoEnv(gym.Env, abc.ABC):
    method __init__ (line 17) | def __init__(
    method build_mjcf_model (line 53) | def build_mjcf_model(self) -> mjcf.RootElement:
    method modify_mjcf_model (line 61) | def modify_mjcf_model(self, mjcf_model: mjcf.RootElement) -> mjcf.Root...
    method initialize_episode (line 77) | def initialize_episode(self) -> None:
    method compute_observation (line 82) | def compute_observation(self) -> Any:
    method compute_reward (line 91) | def compute_reward(self) -> SupportsFloat:
    method set_control (line 95) | def set_control(self, action) -> None:
    method post_compilation (line 103) | def post_compilation(self) -> None:
    method terminate_episode (line 111) | def terminate_episode(self) -> bool:
    method truncate_episode (line 118) | def truncate_episode(self) -> bool:
    method get_reset_info (line 125) | def get_reset_info(self) -> dict:
    method get_step_info (line 129) | def get_step_info(self) -> dict:
    method pre_step (line 133) | def pre_step(self) -> None:
    method post_step (line 140) | def post_step(self) -> None:
    method compile_model_and_data (line 148) | def compile_model_and_data(self):
    method mark_dirty (line 180) | def mark_dirty(self):
    method reset (line 184) | def reset(self, seed: int = None, options=None, **kwargs):
    method set_state (line 212) | def set_state(self, qpos, qvel):
    method step (line 221) | def step(self, action):
    method action_space (line 246) | def action_space(self):
    method observation_space (line 263) | def observation_space(self):
    method set_timesteps (line 270) | def set_timesteps(self, physics_timestep: float, control_timestep: flo...
    method model (line 292) | def model(self) -> mujoco.MjModel:
    method data (line 299) | def data(self) -> mujoco.MjData:
    method mjcf_model (line 306) | def mjcf_model(self) -> mjcf.RootElement:
    method physics_timestep (line 312) | def physics_timestep(self) -> float:
    method control_timestep (line 316) | def control_timestep(self) -> float:
    method launch_passive_viewer (line 322) | def launch_passive_viewer(self, *args, **kwargs):
    method sync_passive_viewer (line 337) | def sync_passive_viewer(self):
    method close_passive_viewer (line 343) | def close_passive_viewer(self):
    method passive_viewer (line 350) | def passive_viewer(self, *args, **kwargs):
    method _initialize_renderer (line 362) | def _initialize_renderer(self):
    method render (line 369) | def render(

FILE: ogbench/manipspace/envs/manipspace_env.py
  class ManipSpaceEnv (line 13) | class ManipSpaceEnv(CustomMuJoCoEnv):
    method __init__ (line 23) | def __init__(
    method observation_space (line 133) | def observation_space(self):
    method action_space (line 145) | def action_space(self):
    method normalize_action (line 153) | def normalize_action(self, action):
    method unnormalize_action (line 158) | def unnormalize_action(self, action):
    method set_tasks (line 162) | def set_tasks(self):
    method build_mjcf_model (line 165) | def build_mjcf_model(self):
    method add_objects (line 245) | def add_objects(self, arena_mjcf):
    method post_compilation (line 248) | def post_compilation(self):
    method post_compilation_objects (line 279) | def post_compilation_objects(self):
    method reset (line 282) | def reset(self, options=None, *args, **kwargs):
    method step (line 314) | def step(self, action):
    method initialize_arm (line 343) | def initialize_arm(self):
    method initialize_episode (line 363) | def initialize_episode(self):
    method set_new_target (line 366) | def set_new_target(self, return_info=True):
    method set_control (line 369) | def set_control(self, action):
    method pre_step (line 419) | def pre_step(self):
    method compute_ob_info (line 424) | def compute_ob_info(self):
    method add_object_info (line 453) | def add_object_info(self, ob_info):
    method get_pixel_observation (line 456) | def get_pixel_observation(self):
    method compute_observation (line 460) | def compute_observation(self):
    method compute_reward (line 481) | def compute_reward(self):
    method get_reset_info (line 484) | def get_reset_info(self):
    method get_step_info (line 492) | def get_step_info(self):
    method terminate_episode (line 496) | def terminate_episode(self):
    method render (line 502) | def render(

FILE: ogbench/manipspace/envs/puzzle_env.py
  class PuzzleEnv (line 8) | class PuzzleEnv(ManipSpaceEnv):
    method __init__ (line 20) | def __init__(self, env_type, *args, **kwargs):
    method set_state (line 61) | def set_state(self, qpos, qvel, button_states):
    method set_tasks (line 66) | def set_tasks(self):
    method add_objects (line 461) | def add_objects(self, arena_mjcf):
    method post_compilation_objects (line 499) | def post_compilation_objects(self):
    method _apply_button_states (line 507) | def _apply_button_states(self):
    method initialize_episode (line 517) | def initialize_episode(self):
    method set_new_target (line 578) | def set_new_target(self, return_info=True, p_stack=0.5):
    method pre_step (line 598) | def pre_step(self):
    method _compute_successes (line 602) | def _compute_successes(self):
    method post_step (line 610) | def post_step(self):
    method add_object_info (line 633) | def add_object_info(self, ob_info):
    method compute_observation (line 653) | def compute_observation(self):
    method compute_oracle_observation (line 684) | def compute_oracle_observation(self):
    method compute_reward (line 688) | def compute_reward(self):

FILE: ogbench/manipspace/envs/scene_env.py
  class SceneEnv (line 9) | class SceneEnv(ManipSpaceEnv):
    method __init__ (line 20) | def __init__(self, env_type, permute_blocks=True, *args, **kwargs):
    method set_state (line 57) | def set_state(self, qpos, qvel, button_states):
    method set_tasks (line 62) | def set_tasks(self):
    method add_objects (line 144) | def add_objects(self, arena_mjcf):
    method post_compilation_objects (line 180) | def post_compilation_objects(self):
    method _apply_button_states (line 207) | def _apply_button_states(self):
    method initialize_episode (line 237) | def initialize_episode(self):
    method _is_in_drawer (line 356) | def _is_in_drawer(self, obj_pos):
    method set_new_target (line 363) | def set_new_target(self, return_info=True, p_stack=0.5):
    method pre_step (line 481) | def pre_step(self):
    method _compute_successes (line 485) | def _compute_successes(self):
    method post_step (line 503) | def post_step(self):
    method add_object_info (line 570) | def add_object_info(self, ob_info):
    method compute_observation (line 630) | def compute_observation(self):
    method compute_oracle_observation (line 680) | def compute_oracle_observation(self):
    method compute_reward (line 701) | def compute_reward(self):

FILE: ogbench/manipspace/lie/se3.py
  class SE3 (line 15) | class SE3:
    method __repr__ (line 28) | def __repr__(self) -> str:
    method identity (line 34) | def identity() -> SE3:
    method from_rotation_and_translation (line 38) | def from_rotation_and_translation(
    method from_matrix (line 46) | def from_matrix(matrix: np.ndarray) -> SE3:
    method sample_uniform (line 54) | def sample_uniform() -> SE3:
    method rotation (line 60) | def rotation(self) -> SO3:
    method translation (line 63) | def translation(self) -> np.ndarray:
    method as_matrix (line 66) | def as_matrix(self) -> np.ndarray:
    method exp (line 73) | def exp(tangent: np.ndarray) -> SE3:
    method log (line 94) | def log(self) -> np.ndarray:
    method adjoint (line 114) | def adjoint(self) -> np.ndarray:
    method inverse (line 123) | def inverse(self) -> SE3:
    method normalize (line 130) | def normalize(self) -> SE3:
    method apply (line 136) | def apply(self, target: np.ndarray) -> np.ndarray:
    method multiply (line 140) | def multiply(self, other: SE3) -> SE3:
    method __matmul__ (line 146) | def __matmul__(self, other: Any) -> Any:

FILE: ogbench/manipspace/lie/so3.py
  class RollPitchYaw (line 16) | class RollPitchYaw:
  class SO3 (line 23) | class SO3:
    method __post_init__ (line 36) | def __post_init__(self) -> None:
    method __repr__ (line 40) | def __repr__(self) -> str:
    method copy (line 44) | def copy(self) -> SO3:
    method from_x_radians (line 48) | def from_x_radians(theta: float) -> SO3:
    method from_y_radians (line 52) | def from_y_radians(theta: float) -> SO3:
    method from_z_radians (line 56) | def from_z_radians(theta: float) -> SO3:
    method from_rpy_radians (line 60) | def from_rpy_radians(
    method from_matrix (line 68) | def from_matrix(matrix: np.ndarray) -> SO3:
    method identity (line 75) | def identity() -> SO3:
    method sample_uniform (line 79) | def sample_uniform() -> SO3:
    method as_matrix (line 97) | def as_matrix(self) -> np.ndarray:
    method compute_roll_radians (line 102) | def compute_roll_radians(self) -> float:
    method compute_pitch_radians (line 106) | def compute_pitch_radians(self) -> float:
    method compute_yaw_radians (line 110) | def compute_yaw_radians(self) -> float:
    method as_rpy_radians (line 114) | def as_rpy_radians(self) -> RollPitchYaw:
    method exp (line 122) | def exp(tangent: np.ndarray) -> SO3:
    method log (line 138) | def log(self) -> np.ndarray:
    method adjoint (line 155) | def adjoint(self) -> np.ndarray:
    method inverse (line 158) | def inverse(self) -> SO3:
    method normalize (line 161) | def normalize(self) -> SO3:
    method apply (line 164) | def apply(self, target: np.ndarray) -> np.ndarray:
    method multiply (line 169) | def multiply(self, other: SO3) -> SO3:
    method __matmul__ (line 184) | def __matmul__(self, other: Any) -> Any:

FILE: ogbench/manipspace/lie/utils.py
  function get_epsilon (line 5) | def get_epsilon(dtype: np.dtype) -> float:
  function skew (line 12) | def skew(x: np.ndarray) -> np.ndarray:
  function mat2quat (line 24) | def mat2quat(mat: np.ndarray):
  function interpolate (line 32) | def interpolate(p0, p1, alpha=0.5):

FILE: ogbench/manipspace/mjcf_utils.py
  function attach (line 9) | def attach(
  function to_string (line 52) | def to_string(
  function get_assets (line 98) | def get_assets(root: mjcf.RootElement) -> dict:
  function safe_find_all (line 106) | def safe_find_all(root: mjcf.RootElement, namespace: str, *args, **kwargs):
  function safe_find (line 114) | def safe_find(root: mjcf.RootElement, namespace: str, identifier: str):
  function add_bounding_box_site (line 122) | def add_bounding_box_site(body: mjcf.Element, lower: np.ndarray, upper: ...

FILE: ogbench/manipspace/oracles/markov/button_markov.py
  class ButtonMarkovOracle (line 6) | class ButtonMarkovOracle(MarkovOracle):
    method __init__ (line 7) | def __init__(self, max_step=100, gripper_always_closed=False, *args, *...
    method reset (line 12) | def reset(self, ob, info):
    method select_action (line 18) | def select_action(self, ob, info):

FILE: ogbench/manipspace/oracles/markov/cube_markov.py
  class CubeMarkovOracle (line 6) | class CubeMarkovOracle(MarkovOracle):
    method __init__ (line 7) | def __init__(self, max_step=200, *args, **kwargs):
    method reset (line 11) | def reset(self, ob, info):
    method select_action (line 18) | def select_action(self, ob, info):

FILE: ogbench/manipspace/oracles/markov/drawer_markov.py
  class DrawerMarkovOracle (line 6) | class DrawerMarkovOracle(MarkovOracle):
    method __init__ (line 7) | def __init__(self, max_step=75, *args, **kwargs):
    method reset (line 11) | def reset(self, ob, info):
    method select_action (line 17) | def select_action(self, ob, info):

FILE: ogbench/manipspace/oracles/markov/markov_oracle.py
  class MarkovOracle (line 4) | class MarkovOracle:
    method __init__ (line 7) | def __init__(self, env, min_norm=0.4):
    method shape_diff (line 23) | def shape_diff(self, diff):
    method shortest_yaw (line 31) | def shortest_yaw(self, eff_yaw, obj_yaw, n=4):
    method print_phase (line 37) | def print_phase(self, phase):
    method done (line 43) | def done(self):
    method reset (line 46) | def reset(self, ob, info):
    method select_action (line 49) | def select_action(self, ob, info):

FILE: ogbench/manipspace/oracles/markov/window_markov.py
  class WindowMarkovOracle (line 6) | class WindowMarkovOracle(MarkovOracle):
    method __init__ (line 7) | def __init__(self, max_step=75, *args, **kwargs):
    method reset (line 11) | def reset(self, ob, info):
    method select_action (line 19) | def select_action(self, ob, info):

FILE: ogbench/manipspace/oracles/plan/button_plan.py
  class ButtonPlanOracle (line 6) | class ButtonPlanOracle(PlanOracle):
    method __init__ (line 7) | def __init__(self, gripper_always_closed=False, *args, **kwargs):
    method compute_keyframes (line 11) | def compute_keyframes(self, plan_input):
    method reset (line 46) | def reset(self, ob, info):

FILE: ogbench/manipspace/oracles/plan/cube_plan.py
  class CubePlanOracle (line 7) | class CubePlanOracle(PlanOracle):
    method __init__ (line 8) | def __init__(
    method compute_keyframes (line 15) | def compute_keyframes(self, plan_input):
    method reset (line 78) | def reset(self, ob, info):

FILE: ogbench/manipspace/oracles/plan/drawer_plan.py
  class DrawerPlanOracle (line 6) | class DrawerPlanOracle(PlanOracle):
    method __init__ (line 7) | def __init__(self, *args, **kwargs):
    method compute_keyframes (line 10) | def compute_keyframes(self, plan_input):
    method reset (line 58) | def reset(self, ob, info):

FILE: ogbench/manipspace/oracles/plan/plan_oracle.py
  class PlanOracle (line 8) | class PlanOracle:
    method __init__ (line 15) | def __init__(self, env, segment_dt=0.4, noise=0.1, noise_smoothing=0.5):
    method above (line 35) | def above(self, pose, z):
    method to_pose (line 44) | def to_pose(self, pos, yaw):
    method get_yaw (line 50) | def get_yaw(self, pose):
    method shortest_yaw (line 56) | def shortest_yaw(self, eff_yaw, obj_yaw, translation, n=4):
    method compute_plan (line 65) | def compute_plan(self, times, poses, grasps):
    method done (line 104) | def done(self):
    method reset (line 107) | def reset(self, ob, info):
    method select_action (line 110) | def select_action(self, ob, info):

FILE: ogbench/manipspace/oracles/plan/window_plan.py
  class WindowPlanOracle (line 6) | class WindowPlanOracle(PlanOracle):
    method __init__ (line 7) | def __init__(self, *args, **kwargs):
    method compute_keyframes (line 10) | def compute_keyframes(self, plan_input):
    method reset (line 58) | def reset(self, ob, info):

FILE: ogbench/manipspace/viewer_utils.py
  class KeyCallback (line 7) | class KeyCallback:
    method __call__ (line 11) | def __call__(self, key: int) -> None:

FILE: ogbench/online_locomotion/ant.py
  class AntEnv (line 14) | class AntEnv(MujocoEnv, utils.EzPickle):
    method __init__ (line 30) | def __init__(
    method healthy_reward (line 94) | def healthy_reward(self):
    method control_cost (line 97) | def control_cost(self, action):
    method contact_forces (line 102) | def contact_forces(self):
    method contact_cost (line 109) | def contact_cost(self):
    method is_healthy (line 114) | def is_healthy(self):
    method terminated (line 121) | def terminated(self):
    method step (line 125) | def step(self, action):
    method _get_obs (line 164) | def _get_obs(self):
    method reset_model (line 177) | def reset_model(self):

FILE: ogbench/online_locomotion/ant_ball.py
  class AntBallEnv (line 10) | class AntBallEnv(AntEnv):
    method __init__ (line 13) | def __init__(self, xml_file=None, *args, **kwargs):
    method reset (line 56) | def reset(self, options=None, *args, **kwargs):
    method step (line 69) | def step(self, action):
    method set_goal (line 92) | def set_goal(self, goal_xy):
    method get_agent_ball_xy (line 96) | def get_agent_ball_xy(self):
    method set_agent_ball_xy (line 102) | def set_agent_ball_xy(self, agent_xy, ball_xy):
    method _get_obs (line 109) | def _get_obs(self):

FILE: ogbench/online_locomotion/humanoid.py
  function _sigmoids (line 18) | def _sigmoids(x, value_at_1, sigmoid):
  function tolerance (line 68) | def tolerance(x, bounds=(0.0, 0.0), margin=0.0, sigmoid='gaussian', valu...
  class HumanoidEnv (line 85) | class HumanoidEnv(MujocoEnv, utils.EzPickle):
    method __init__ (line 100) | def __init__(
    method step (line 130) | def step(self, action):
    method _step_mujoco_simulation (line 149) | def _step_mujoco_simulation(self, ctrl, n_frames):
    method _get_obs (line 162) | def _get_obs(self):
    method _get_reward (line 180) | def _get_reward(self):
    method disable (line 208) | def disable(self, *flags):
    method reset_model (line 224) | def reset_model(self):
    method get_xy (line 243) | def get_xy(self):
    method set_xy (line 246) | def set_xy(self, xy):

FILE: ogbench/online_locomotion/wrappers.py
  class GymXYWrapper (line 6) | class GymXYWrapper(gymnasium.Wrapper):
    method __init__ (line 9) | def __init__(self, env, resample_interval=100):
    method reset (line 25) | def reset(self, *args, **kwargs):
    method step (line 33) | def step(self, action):
  class DMCHumanoidXYWrapper (line 51) | class DMCHumanoidXYWrapper(GymXYWrapper):
    method step (line 54) | def step(self, action):

FILE: ogbench/powderworld/behaviors.py
  class Behavior (line 4) | class Behavior:
    method __init__ (line 7) | def __init__(self, env):
    method done (line 18) | def done(self):
    method reset (line 21) | def reset(self, ob, info):
    method select_action (line 24) | def select_action(self, ob, info):
  class FillBehavior (line 34) | class FillBehavior(Behavior):
    method reset (line 37) | def reset(self, ob, info):
  class LineBehavior (line 61) | class LineBehavior(Behavior):
    method reset (line 64) | def reset(self, ob, info):
  class SquareBehavior (line 85) | class SquareBehavior(Behavior):
    method reset (line 88) | def reset(self, ob, info):

FILE: ogbench/powderworld/powderworld_env.py
  class PowderworldEnv (line 9) | class PowderworldEnv(gymnasium.Env):
    method __init__ (line 21) | def __init__(
    method set_tasks (line 91) | def set_tasks(self):
    method reset (line 284) | def reset(self, *, seed=None, options=None):
    method step (line 354) | def step(self, action):
    method semantic_action_to_action (line 429) | def semantic_action_to_action(self, elem_name, x, y):
    method sample_action (line 439) | def sample_action(self):
    method sample_semantic_action (line 446) | def sample_semantic_action(self):
    method render (line 453) | def render(self):
    method _get_ob (line 462) | def _get_ob(self):

FILE: ogbench/powderworld/sim.py
  function get_below (line 70) | def get_below(x):
  function get_above (line 74) | def get_above(x):
  function get_left (line 78) | def get_left(x):
  function get_right (line 82) | def get_right(x):
  function get_in_cardinal_direction (line 86) | def get_in_cardinal_direction(x, directions):
  function interp (line 94) | def interp(switch, if_false, if_true):
  function interp_int (line 98) | def interp_int(switch, if_false, if_true: int):
  function interp2 (line 102) | def interp2(switch_a, switch_b, if_false, if_a, if_b):
  function interp_swaps8 (line 106) | def interp_swaps8(swaps, world, w0, w1, w2, w3, w4, w5, w6, w7):
  function interp_swaps4 (line 119) | def interp_swaps4(swaps, world, w0, w1, w2, w3):
  function normalize (line 128) | def normalize(x, p=2, axis=0, eps=1e-12):
  function conv2d (line 134) | def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, gr...
  function conv2d_pad (line 204) | def conv2d_pad(input, padding, padding_mode, stride, dilation, kernel_si...
  class PWSim (line 242) | class PWSim:
    method __init__ (line 243) | def __init__(self):
    method register_update_rules (line 284) | def register_update_rules(self):
    method add_element (line 311) | def add_element(self, world_slice, element_name, wind=None):
    method add_element_rc (line 322) | def add_element_rc(self, world_slice, rr, cc, element_name):
    method id_to_pw (line 327) | def id_to_pw(self, world_ids):
    method np_to_pw (line 332) | def np_to_pw(self, np_world):
    method get_elem (line 337) | def get_elem(self, world, elemname):
    method get_bool (line 341) | def get_bool(self, world, elemname):
    method direction_func (line 345) | def direction_func(self, d, x):
    method forward (line 363) | def forward(self, world):
  class PWRenderer (line 386) | class PWRenderer:
    method __init__ (line 387) | def __init__(self):
    method forward (line 425) | def forward(self, world):
    method render (line 449) | def render(self, world):
  class BehaviorGravity (line 461) | class BehaviorGravity:
    method __init__ (line 470) | def __init__(self, pw):
    method check_filter (line 473) | def check_filter(self, world):
    method forward (line 476) | def forward(self, world, info):
  class BehaviorSand (line 504) | class BehaviorSand:
    method __init__ (line 512) | def __init__(self, pw):
    method check_filter (line 515) | def check_filter(self, world):
    method forward (line 518) | def forward(self, world, info):
  class BehaviorStone (line 574) | class BehaviorStone:
    method __init__ (line 577) | def __init__(self, pw):
    method check_filter (line 583) | def check_filter(self, world):
    method forward (line 586) | def forward(self, world, info):
  class BehaviorFluidFlow (line 593) | class BehaviorFluidFlow:
    method __init__ (line 598) | def __init__(self, pw):
    method check_filter (line 601) | def check_filter(self, world):
    method forward (line 604) | def forward(self, world, info):
  class BehaviorIce (line 670) | class BehaviorIce:
    method __init__ (line 675) | def __init__(self, pw):
    method check_filter (line 678) | def check_filter(self, world):
    method forward (line 681) | def forward(self, world, info):
  class BehaviorWater (line 696) | class BehaviorWater:
    method __init__ (line 701) | def __init__(self, pw):
    method check_filter (line 704) | def check_filter(self, world):
    method forward (line 707) | def forward(self, world, info):
  class BehaviorFire (line 716) | class BehaviorFire:
    method __init__ (line 721) | def __init__(self, pw):
    method check_filter (line 724) | def check_filter(self, world):
    method forward (line 727) | def forward(self, world, info):
  class BehaviorPlant (line 790) | class BehaviorPlant:
    method __init__ (line 795) | def __init__(self, pw):
    method check_filter (line 798) | def check_filter(self, world):
    method forward (line 801) | def forward(self, world, info):
  class BehaviorLava (line 826) | class BehaviorLava:
    method __init__ (line 831) | def __init__(self, pw):
    method check_filter (line 834) | def check_filter(self, world):
    method forward (line 837) | def forward(self, world, info):
  class BehaviorAcid (line 848) | class BehaviorAcid:
    method __init__ (line 853) | def __init__(self, pw):
    method check_filter (line 856) | def check_filter(self, world):
    method forward (line 859) | def forward(self, world, info):
  class BehaviorCloner (line 880) | class BehaviorCloner:
    method __init__ (line 885) | def __init__(self, pw):
    method check_filter (line 888) | def check_filter(self, world):
    method forward (line 891) | def forward(self, world, info):
  class BehaviorVelocity (line 919) | class BehaviorVelocity:
    method __init__ (line 924) | def __init__(self, pw):
    method check_filter (line 927) | def check_filter(self, world):
    method forward (line 930) | def forward(self, world, info):
  class BehaviorFish (line 985) | class BehaviorFish:
    method __init__ (line 992) | def __init__(self, pw):
    method check_filter (line 995) | def check_filter(self, world):
    method forward (line 998) | def forward(self, world, info):
  class BehaviorBird (line 1031) | class BehaviorBird:
    method __init__ (line 1036) | def __init__(self, pw):
    method check_filter (line 1050) | def check_filter(self, world):
    method forward (line 1053) | def forward(self, world, info):
  class BehaviorKangaroo (line 1078) | class BehaviorKangaroo:
    method __init__ (line 1083) | def __init__(self, pw):
    method check_filter (line 1086) | def check_filter(self, world):
    method forward (line 1089) | def forward(self, world, info):
  class BehaviorMole (line 1109) | class BehaviorMole:
    method __init__ (line 1114) | def __init__(self, pw):
    method check_filter (line 1117) | def check_filter(self, world):
    method forward (line 1120) | def forward(self, world, info):
  class BehaviorLemming (line 1165) | class BehaviorLemming:
    method __init__ (line 1170) | def __init__(self, pw):
    method check_filter (line 1173) | def check_filter(self, world):
    method forward (line 1176) | def forward(self, world, info):
  class BehaviorSnake (line 1214) | class BehaviorSnake:
    method __init__ (line 1224) | def __init__(self, pw):
    method check_filter (line 1227) | def check_filter(self, world):
    method forward (line 1230) | def forward(self, world, info):

FILE: ogbench/relabel_utils.py
  function relabel_dataset (line 4) | def relabel_dataset(env_name, env, dataset):
  function add_oracle_reps (line 93) | def add_oracle_reps(env_name, env, dataset):

FILE: ogbench/utils.py
  function load_dataset (line 14) | def load_dataset(dataset_path, ob_dtype=np.float32, action_dtype=np.floa...
  function download_datasets (line 99) | def download_datasets(dataset_names, dataset_dir=DEFAULT_DATASET_DIR):
  function make_env_and_datasets (line 134) | def make_env_and_datasets(