SYMBOL INDEX (602 symbols across 36 files)

FILE: eval.py
  class JointsTrajectoryActionMode (line 41) | class JointsTrajectoryActionMode(ArmActionMode):
    method __init__ (line 44) | def __init__(self, points: int):
    method action (line 47) | def action(self, scene: Scene, action: np.ndarray, ignore_collisions: ...
    method _pre_proc_traj (line 68) | def _pre_proc_traj(self, action):
    method action_shape (line 80) | def action_shape(self, scene: Scene) -> tuple:
    method set_callable_each_step (line 83) | def set_callable_each_step(self, callable_each_step):
    method record_end (line 86) | def record_end(self, scene, steps=60, step_scene=True):
  function eval_seed (line 93) | def eval_seed(
  function main (line 242) | def main(eval_cfg: DictConfig) -> None:

FILE: extra_scripts/dataset_generator.py
  function check_and_make (line 52) | def check_and_make(dir):
  function save_demo (line 57) | def save_demo(demo, example_path):
  class CustomMoveArmThenGripper (line 172) | class CustomMoveArmThenGripper(MoveArmThenGripper):
    method action_bounds (line 173) | def action_bounds(self):
  class AbsCustomMoveArmThenGripper (line 180) | class AbsCustomMoveArmThenGripper(MoveArmThenGripper):
    method action_bounds (line 202) | def action_bounds(self):
  function run (line 209) | def run(
  function main (line 410) | def main(argv):

FILE: helpers/clip/core/attention.py
  class Attention (line 11) | class Attention(nn.Module):
    method __init__ (line 14) | def __init__(self, stream_fcn, in_shape, n_rotations, preprocess, cfg,...
    method _build_nets (line 37) | def _build_nets(self):
    method attend (line 44) | def attend(self, x):
    method forward (line 47) | def forward(self, inp_img, softmax=True):

FILE: helpers/clip/core/attention_image_goal.py
  class AttentionImageGoal (line 9) | class AttentionImageGoal(Attention):
    method __init__ (line 12) | def __init__(self, stream_fcn, in_shape, n_rotations, preprocess, cfg,...
    method forward (line 15) | def forward(self, inp_img, goal_img, softmax=True):

FILE: helpers/clip/core/clip.py
  class Bottleneck (line 41) | class Bottleneck(nn.Module):
    method __init__ (line 44) | def __init__(self, inplanes, planes, stride=1):
    method forward (line 84) | def forward(self, x: torch.Tensor):
  class AttentionPool2d (line 100) | class AttentionPool2d(nn.Module):
    method __init__ (line 101) | def __init__(
    method forward (line 114) | def forward(self, x):
  class ModifiedResNet (line 147) | class ModifiedResNet(nn.Module):
    method __init__ (line 155) | def __init__(self, layers, output_dim, heads, input_resolution=224, wi...
    method _make_layer (line 186) | def _make_layer(self, planes, blocks, stride=1):
    method forward (line 195) | def forward(self, x):
    method prepool (line 200) | def prepool(self, x):
    method prepool_im (line 219) | def prepool_im(self, x):
  class LayerNorm (line 245) | class LayerNorm(nn.LayerNorm):
    method forward (line 248) | def forward(self, x: torch.Tensor):
  class QuickGELU (line 254) | class QuickGELU(nn.Module):
    method forward (line 255) | def forward(self, x: torch.Tensor):
  class ResidualAttentionBlock (line 259) | class ResidualAttentionBlock(nn.Module):
    method __init__ (line 260) | def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor ...
    method attention (line 277) | def attention(self, x: torch.Tensor):
    method forward (line 285) | def forward(self, x: torch.Tensor):
  class Transformer (line 291) | class Transformer(nn.Module):
    method __init__ (line 292) | def __init__(
    method forward (line 302) | def forward(self, x: torch.Tensor):
  class VisualTransformer (line 306) | class VisualTransformer(nn.Module):
    method __init__ (line 307) | def __init__(
    method forward (line 339) | def forward(self, x: torch.Tensor):
    method forward_spatial (line 367) | def forward_spatial(self, x: torch.Tensor):
  class CLIP (line 392) | class CLIP(nn.Module):
    method __init__ (line 393) | def __init__(
    method initialize_parameters (line 451) | def initialize_parameters(self):
    method build_attention_mask (line 487) | def build_attention_mask(self):
    method dtype (line 496) | def dtype(self):
    method encode_image (line 499) | def encode_image(self, image):
    method encode_text (line 502) | def encode_text(self, text):
    method encode_text_with_embeddings (line 517) | def encode_text_with_embeddings(self, text):
    method forward (line 533) | def forward(self, image, text):
  function convert_weights (line 550) | def convert_weights(model: nn.Module):
  function build_model (line 579) | def build_model(state_dict: dict):
  function _download (line 654) | def _download(url: str, root: str = os.path.expanduser("~/.cache/clip")):
  function available_models (line 696) | def available_models():
  function load_clip (line 700) | def load_clip(
  function tokenize (line 791) | def tokenize(texts: Union[str, List[str]], context_length: int = 77):

FILE: helpers/clip/core/fusion.py
  class DotAttn (line 7) | class DotAttn(nn.Module):
    method forward (line 10) | def forward(self, inp, h):
    method softmax (line 14) | def softmax(self, inp, h):
  class ScaledDotAttn (line 20) | class ScaledDotAttn(nn.Module):
    method forward (line 23) | def forward(self, inp, h):
    method softmax (line 27) | def softmax(self, inp, h):
  class Fusion (line 33) | class Fusion(nn.Module):
    method __init__ (line 36) | def __init__(self, input_dim=3):
    method tile_x2 (line 40) | def tile_x2(self, x1, x2, x2_proj=None):
    method forward (line 48) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionAdd (line 52) | class FusionAdd(Fusion):
    method __init__ (line 55) | def __init__(self, input_dim=3):
    method forward (line 58) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionMult (line 64) | class FusionMult(Fusion):
    method __init__ (line 67) | def __init__(self, input_dim=3):
    method forward (line 70) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionMax (line 76) | class FusionMax(Fusion):
    method __init__ (line 79) | def __init__(self, input_dim=3):
    method forward (line 82) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionConcat (line 88) | class FusionConcat(Fusion):
    method __init__ (line 91) | def __init__(self, input_dim=3):
    method forward (line 94) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionConv (line 100) | class FusionConv(Fusion):
    method __init__ (line 103) | def __init__(self, input_dim=3):
    method forward (line 110) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionConvLat (line 118) | class FusionConvLat(Fusion):
    method __init__ (line 121) | def __init__(self, input_dim=3, output_dim=3):
    method forward (line 127) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionFiLM (line 142) | class FusionFiLM(Fusion):
    method __init__ (line 148) | def __init__(self, input_dim=3, output_dim=3):
    method forward (line 151) | def forward(self, x1, x2, gamma, beta):
  class FusionDeepConv (line 157) | class FusionDeepConv(Fusion):
    method __init__ (line 160) | def __init__(self, input_dim=3):
    method forward (line 171) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionMultWord (line 179) | class FusionMultWord(nn.Module):
    method __init__ (line 182) | def __init__(self, input_dim=3):
    method forward (line 186) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionWordAttention (line 199) | class FusionWordAttention(nn.Module):
    method __init__ (line 202) | def __init__(self, input_dim=3):
    method forward (line 207) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class FusionSentenceAttention (line 226) | class FusionSentenceAttention(nn.Module):
    method __init__ (line 229) | def __init__(self, input_dim=3):
    method forward (line 234) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):
  class CrossModalAttention2d (line 248) | class CrossModalAttention2d(nn.Module):
    method __init__ (line 251) | def __init__(
    method forward (line 277) | def forward(self, x, l, l_mask):
  class FusionMultiHeadedWordAttention (line 329) | class FusionMultiHeadedWordAttention(nn.Module):
    method __init__ (line 332) | def __init__(self, input_dim=3):
    method forward (line 351) | def forward(self, x1, x2, x2_mask=None, x2_proj=None):

FILE: helpers/clip/core/resnet.py
  class IdentityBlock (line 6) | class IdentityBlock(nn.Module):
    method __init__ (line 7) | def __init__(
    method forward (line 30) | def forward(self, x):
  class ConvBlock (line 40) | class ConvBlock(nn.Module):
    method __init__ (line 41) | def __init__(
    method forward (line 69) | def forward(self, x):
  class ResNet43_8s (line 79) | class ResNet43_8s(nn.Module):
    method __init__ (line 80) | def __init__(self, input_shape, output_dim, cfg, device, preprocess):
    method _make_layers (line 92) | def _make_layers(self):
    method forward (line 165) | def forward(self, x):

FILE: helpers/clip/core/simple_tokenizer.py
  function default_bpe (line 11) | def default_bpe():
  function bytes_to_unicode (line 18) | def bytes_to_unicode():
  function get_pairs (line 44) | def get_pairs(word):
  function basic_clean (line 56) | def basic_clean(text):
  function whitespace_clean (line 62) | def whitespace_clean(text):
  class SimpleTokenizer (line 68) | class SimpleTokenizer(object):
    method __init__ (line 69) | def __init__(self, bpe_path: str = default_bpe()):
    method bpe (line 92) | def bpe(self, token):
    method encode (line 133) | def encode(self, text):
    method decode (line 143) | def decode(self, tokens):

FILE: helpers/clip/core/transport.py
  class Transport (line 9) | class Transport(nn.Module):
    method __init__ (line 10) | def __init__(
    method _build_nets (line 45) | def _build_nets(self):
    method correlate (line 54) | def correlate(self, in0, in1, softmax):
    method transport (line 70) | def transport(self, in_tensor, crop):
    method forward (line 75) | def forward(self, inp_img, p, softmax=True):

FILE: helpers/clip/core/transport_image_goal.py
  class TransportImageGoal (line 9) | class TransportImageGoal(nn.Module):
    method __init__ (line 12) | def __init__(
    method _build_nets (line 53) | def _build_nets(self):
    method correlate (line 61) | def correlate(self, in0, in1, softmax):
    method forward (line 77) | def forward(self, inp_img, goal_img, p, softmax=True):

FILE: helpers/clip/core/unet.py
  class DoubleConv (line 8) | class DoubleConv(nn.Module):
    method __init__ (line 11) | def __init__(self, in_channels, out_channels, mid_channels=None):
    method forward (line 24) | def forward(self, x):
  class Down (line 28) | class Down(nn.Module):
    method __init__ (line 31) | def __init__(self, in_channels, out_channels):
    method forward (line 37) | def forward(self, x):
  class Up (line 41) | class Up(nn.Module):
    method __init__ (line 44) | def __init__(self, in_channels, out_channels, bilinear=True):
    method forward (line 57) | def forward(self, x1, x2):
  class OutConv (line 71) | class OutConv(nn.Module):
    method __init__ (line 72) | def __init__(self, in_channels, out_channels):
    method forward (line 76) | def forward(self, x):

FILE: helpers/custom_rlbench_env.py
  class CustomRLBenchEnv (line 18) | class CustomRLBenchEnv(RLBenchEnv):
    method __init__ (line 19) | def __init__(
    method observation_elements (line 60) | def observation_elements(self) -> List[ObservationElement]:
    method extract_obs (line 70) | def extract_obs(self, obs: Observation, t=None, prev_action=None):
    method launch (line 107) | def launch(self):
    method reset (line 119) | def reset(self) -> dict:
    method register_callback (line 132) | def register_callback(self, func):
    method _my_callback (line 135) | def _my_callback(self):
    method _append_final_frame (line 141) | def _append_final_frame(self, success: bool):
    method step (line 150) | def step(self, act_result: ActResult) -> Transition:
    method _extract_traj_info (line 205) | def _extract_traj_info(self, demo):
    method reset_to_demo (line 231) | def reset_to_demo(self, i):
  class CustomMultiTaskRLBenchEnv (line 255) | class CustomMultiTaskRLBenchEnv(MultiTaskRLBenchEnv):
    method __init__ (line 256) | def __init__(
    method observation_elements (line 299) | def observation_elements(self) -> List[ObservationElement]:
    method extract_obs (line 309) | def extract_obs(self, obs: Observation, t=None, prev_action=None):
    method launch (line 341) | def launch(self):
    method reset (line 353) | def reset(self) -> dict:
    method register_callback (line 363) | def register_callback(self, func):
    method _my_callback (line 366) | def _my_callback(self):
    method _append_final_frame (line 372) | def _append_final_frame(self, success: bool):
    method step (line 381) | def step(self, act_result: ActResult) -> Transition:
    method reset_to_demo (line 441) | def reset_to_demo(self, i, variation_number=-1):

FILE: helpers/demo_loading_utils.py
  function _is_stopped (line 8) | def _is_stopped(demo, i, obs, stopped_buffer, delta=0.1):
  function keypoint_discovery (line 25) | def keypoint_discovery(demo: Demo, stopping_delta=0.1) -> List[int]:
  function find_minimum_difference (line 91) | def find_minimum_difference(lst):

FILE: helpers/network_utils.py
  function act_layer (line 14) | def act_layer(act):
  function norm_layer2d (line 29) | def norm_layer2d(norm, channels):
  function norm_layer1d (line 42) | def norm_layer1d(norm, num_channels):
  class FiLMBlock (line 53) | class FiLMBlock(nn.Module):
    method __init__ (line 54) | def __init__(self):
    method forward (line 57) | def forward(self, x, gamma, beta):
  class Conv2DBlock (line 66) | class Conv2DBlock(nn.Module):
    method __init__ (line 67) | def __init__(
    method forward (line 120) | def forward(self, x):
  class Conv2DFiLMBlock (line 127) | class Conv2DFiLMBlock(Conv2DBlock):
    method __init__ (line 128) | def __init__(
    method forward (line 150) | def forward(self, x, gamma, beta):
  class Conv3DBlock (line 158) | class Conv3DBlock(nn.Module):
    method __init__ (line 159) | def __init__(
    method forward (line 210) | def forward(self, x):
  class ConvTranspose3DBlock (line 217) | class ConvTranspose3DBlock(nn.Module):
    method __init__ (line 218) | def __init__(
    method forward (line 268) | def forward(self, x):
  class Conv2DUpsampleBlock (line 275) | class Conv2DUpsampleBlock(nn.Module):
    method __init__ (line 276) | def __init__(
    method forward (line 299) | def forward(self, x):
  class Conv3DUpsampleBlock (line 303) | class Conv3DUpsampleBlock(nn.Module):
    method __init__ (line 304) | def __init__(
    method forward (line 327) | def forward(self, x):
  class DenseBlock (line 331) | class DenseBlock(nn.Module):
    method __init__ (line 332) | def __init__(self, in_features, out_features, norm=None, activation=No...
    method forward (line 364) | def forward(self, x):
  class SiameseNet (line 371) | class SiameseNet(nn.Module):
    method __init__ (line 372) | def __init__(
    method build (line 390) | def build(self):
    method forward (line 411) | def forward(self, x):
  class CNNAndFcsNet (line 421) | class CNNAndFcsNet(nn.Module):
    method __init__ (line 422) | def __init__(
    method build (line 445) | def build(self):
    method forward (line 472) | def forward(self, observations, low_dim_ins):
  class CNNLangAndFcsNet (line 482) | class CNNLangAndFcsNet(nn.Module):
    method __init__ (line 483) | def __init__(
    method build (line 508) | def build(self):
    method forward (line 541) | def forward(self, observations, low_dim_ins, lang_goal_emb):
  function pair (line 566) | def pair(t):
  class PreNorm (line 573) | class PreNorm(nn.Module):
    method __init__ (line 574) | def __init__(self, dim, fn):
    method forward (line 579) | def forward(self, x, **kwargs):
  class FeedForward (line 583) | class FeedForward(nn.Module):
    method __init__ (line 584) | def __init__(self, dim, hidden_dim, dropout=0.0):
    method forward (line 594) | def forward(self, x):
  class Attention (line 598) | class Attention(nn.Module):
    method __init__ (line 599) | def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
    method forward (line 618) | def forward(self, x):
  class Transformer (line 632) | class Transformer(nn.Module):
    method __init__ (line 633) | def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.0):
    method forward (line 651) | def forward(self, x):
  class ViT (line 663) | class ViT(nn.Module):
    method __init__ (line 664) | def __init__(
    method forward (line 712) | def forward(self, img):
  class ViTLangAndFcsNet (line 727) | class ViTLangAndFcsNet(nn.Module):
    method __init__ (line 728) | def __init__(
    method build (line 753) | def build(self):
    method forward (line 785) | def forward(self, observations, low_dim_ins, lang_goal_emb):
  class Conv3DInceptionBlockUpsampleBlock (line 808) | class Conv3DInceptionBlockUpsampleBlock(nn.Module):
    method __init__ (line 809) | def __init__(
    method forward (line 836) | def forward(self, x):
  class Conv3DInceptionBlock (line 840) | class Conv3DInceptionBlock(nn.Module):
    method __init__ (line 841) | def __init__(
    method forward (line 886) | def forward(self, x):
  class ConvTransposeUp3DBlock (line 901) | class ConvTransposeUp3DBlock(nn.Module):
    method __init__ (line 902) | def __init__(
    method forward (line 941) | def forward(self, x):
  class SpatialSoftmax3D (line 948) | class SpatialSoftmax3D(torch.nn.Module):
    method __init__ (line 949) | def __init__(self, depth, height, width, channel):
    method forward (line 974) | def forward(self, feature):

FILE: helpers/optim/lamb.py
  class Lamb (line 27) | class Lamb(Optimizer):
    method __init__ (line 45) | def __init__(
    method step (line 60) | def step(self, closure=None):

FILE: helpers/preprocess_agent.py
  class PreprocessAgent (line 14) | class PreprocessAgent(Agent):
    method __init__ (line 15) | def __init__(self, pose_agent: Agent, norm_rgb: bool = True):
    method build (line 19) | def build(self, training: bool, device: torch.device = None):
    method _norm_rgb_ (line 22) | def _norm_rgb_(self, x):
    method update (line 25) | def update(self, step: int, replay_sample: dict) -> dict:
    method act (line 38) | def act(self, step: int, observation: dict, deterministic=False) -> Ac...
    method update_summaries (line 49) | def update_summaries(self) -> List[Summary]:
    method act_summaries (line 98) | def act_summaries(self) -> List[Summary]:
    method load_weights (line 101) | def load_weights(self, savedir: str):
    method save_weights (line 104) | def save_weights(self, savedir: str):
    method reset (line 107) | def reset(self) -> None:

FILE: helpers/utils.py
  function loss_weights (line 30) | def loss_weights(replay_sample, beta=1.0):
  function soft_updates (line 39) | def soft_updates(net, target_net, tau):
  function stack_on_channel (line 44) | def stack_on_channel(x):
  function normalize_quaternion (line 49) | def normalize_quaternion(quat):
  function correct_rotation_instability (line 53) | def correct_rotation_instability(disc, resolution):
  function check_gimbal_lock (line 68) | def check_gimbal_lock(pred_rot_and_grip, gt_rot_and_grip, resolution):
  function quaternion_to_discrete_euler (line 82) | def quaternion_to_discrete_euler(quaternion, resolution):
  function discrete_euler_to_quaternion (line 90) | def discrete_euler_to_quaternion(discrete_euler, resolution):
  function point_to_voxel_index (line 95) | def point_to_voxel_index(
  function voxel_index_to_point (line 109) | def voxel_index_to_point(
  function point_to_pixel_index (line 117) | def point_to_pixel_index(
  function _compute_initial_camera_pose (line 129) | def _compute_initial_camera_pose(scene):
  function _from_trimesh_scene (line 145) | def _from_trimesh_scene(trimesh_scene, bg_color=None, ambient_light=None):
  function _create_bounding_box (line 160) | def _create_bounding_box(scene, voxel_size, res):
  function create_voxel_scene (line 181) | def create_voxel_scene(
  function visualise_voxel (line 233) | def visualise_voxel(
  function preprocess (line 276) | def preprocess(img, dist="transporter"):
  function rand_dist (line 337) | def rand_dist(size, min=-1.0, max=1.0):
  function rand_discrete (line 341) | def rand_discrete(size, min=0, max=1):
  function split_list (line 347) | def split_list(lst, n):
  function extract_obs (line 352) | def extract_obs(
  function create_obs_config (line 418) | def create_obs_config(
  function get_device (line 459) | def get_device(gpu):

FILE: peract/agents/peract_diffusion/launch_utils.py
  function create_replay (line 40) | def create_replay(
  function _get_action (line 146) | def _get_action(
  function _add_keypoints_to_replay (line 196) | def _add_keypoints_to_replay(
  function fill_replay (line 284) | def fill_replay(
  function fill_multi_task_replay (line 370) | def fill_multi_task_replay(
  function create_diffusion (line 441) | def create_diffusion(cfgs: DictConfig):
  function create_agent (line 450) | def create_agent(cfg: DictConfig):

FILE: peract/agents/peract_diffusion/perceiver_lang_io.py
  function exists (line 24) | def exists(val):
  function default (line 28) | def default(val, d):
  function cache_fn (line 32) | def cache_fn(f):
  function fourier_encode (line 48) | def fourier_encode(x, max_freq, num_bands=4):
  class PreNorm (line 64) | class PreNorm(nn.Module):
    method __init__ (line 65) | def __init__(self, dim, fn, context_dim=None):
    method forward (line 71) | def forward(self, x, **kwargs):
  class GEGLU (line 82) | class GEGLU(nn.Module):
    method forward (line 83) | def forward(self, x):
  class FeedForward (line 88) | class FeedForward(nn.Module):
    method __init__ (line 89) | def __init__(self, dim, mult=4):
    method forward (line 95) | def forward(self, x):
  class Attention (line 99) | class Attention(nn.Module):
    method __init__ (line 100) | def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, ...
    method forward (line 113) | def forward(self, x, context=None, mask=None):
  class PerceiverVoxelLangEncoder (line 142) | class PerceiverVoxelLangEncoder(nn.Module):
    method __init__ (line 143) | def __init__(
    method encode_text (line 394) | def encode_text(self, x):
    method forward (line 403) | def forward(

FILE: peract/agents/peract_diffusion/qattention_peract_bc_agent.py
  class QFunction (line 30) | class QFunction(nn.Module):
    method __init__ (line 31) | def __init__(
    method _argmax_3d (line 50) | def _argmax_3d(self, tensor_orig):
    method choose_highest_action (line 56) | def choose_highest_action(self, q_trans, q_rot_grip, q_collision):
    method forward (line 79) | def forward(
  class QAttentionPerActBCAgent (line 127) | class QAttentionPerActBCAgent(Agent):
    method __init__ (line 128) | def __init__(
    method build (line 196) | def build(self, training: bool, device: torch.device = None):
    method _extract_crop (line 320) | def _extract_crop(self, pixel_action, observation):
    method _preprocess_inputs (line 335) | def _preprocess_inputs(self, replay_sample):
    method _act_preprocess_inputs (line 347) | def _act_preprocess_inputs(self, observation):
    method _get_value_from_voxel_index (line 357) | def _get_value_from_voxel_index(self, q, voxel_idx):
    method _get_value_from_rot_and_grip (line 369) | def _get_value_from_rot_and_grip(self, rot_grip_q, rot_and_grip_idx):
    method _celoss (line 388) | def _celoss(self, pred, labels):
    method _softmax_q_trans (line 391) | def _softmax_q_trans(self, q):
    method _softmax_q_rot_grip (line 395) | def _softmax_q_rot_grip(self, q_rot_grip):
    method _softmax_ignore_collision (line 422) | def _softmax_ignore_collision(self, q_collision):
    method update (line 426) | def update(self, step: int, replay_sample: dict) -> dict:
    method act (line 601) | def act(self, step: int, observation: dict, deterministic=False) -> Ac...
    method update_summaries (line 715) | def update_summaries(self) -> List[Summary]:
    method act_summaries (line 748) | def act_summaries(self) -> List[Summary]:
    method load_weights (line 762) | def load_weights(self, savedir: str):
    method save_weights (line 786) | def save_weights(self, savedir: str):

FILE: peract/agents/peract_diffusion/qattention_stack_agent.py
  class QAttentionStackAgent (line 23) | class QAttentionStackAgent(Agent):
    method __init__ (line 24) | def __init__(
    method build (line 49) | def build(self, training: bool, device=None) -> None:
    method update (line 56) | def update(self, step: int, replay_sample: dict) -> dict:
    method act (line 67) | def act(self, step: int, observation: dict, deterministic=False) -> Ac...
    method _plot (line 152) | def _plot(self, predicted_trajs, pcds, rgbs, gt_trajs=None):
    method _prepare_diffusion_data_dict (line 224) | def _prepare_diffusion_data_dict(
    method update_summaries (line 257) | def update_summaries(self) -> List[Summary]:
    method act_summaries (line 263) | def act_summaries(self) -> List[Summary]:
    method load_weights (line 269) | def load_weights(self, savedir: str, backbone: str = "unet"):
    method save_weights (line 286) | def save_weights(self, savedir: str):

FILE: peract/voxel/augmentation.py
  function _axis_angle_rotation (line 8) | def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
  function standardize_quaternion (line 38) | def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
  function _sqrt_positive_part (line 53) | def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
  function quaternion_to_matrix (line 64) | def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
  function matrix_to_quaternion (line 96) | def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
  function euler_angles_to_matrix (line 158) | def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) ...
  function perturb_se3 (line 187) | def perturb_se3(pcd, trans_shift_4x4, rot_shift_4x4, action_gripper_4x4,...
  function apply_se3_augmentation (line 255) | def apply_se3_augmentation(

FILE: peract/voxel/voxel_grid.py
  class VoxelGrid (line 15) | class VoxelGrid(nn.Module):
    method __init__ (line 16) | def __init__(
    method _broadcast (line 132) | def _broadcast(self, src: torch.Tensor, other: torch.Tensor, dim: int):
    method _scatter_mean (line 143) | def _scatter_mean(
    method _scatter_nd (line 165) | def _scatter_nd(self, indices, updates):
    method coords_to_bounding_voxel_grid (line 191) | def coords_to_bounding_voxel_grid(

FILE: rk_diffuser/dataset/realworld_dataset.py
  class RealWorldDataset (line 20) | class RealWorldDataset:
    method __init__ (line 23) | def __init__(
    method _load_dataset (line 152) | def _load_dataset(self, task, data_raw_path, demo_dict):
    method _downsample_images (line 187) | def _downsample_images(self, demo):
    method _extract_obs_from_demo (line 210) | def _extract_obs_from_demo(self, demo: dict):
    method _keypoint_discovery (line 304) | def _keypoint_discovery(self, obs: dict, stopping_delta=0.1) -> List[i...
    method _is_stopped (line 355) | def _is_stopped(self, joint_poses, i, stopped_buffer, delta=0.1):
    method _is_gripper_open (line 386) | def _is_gripper_open(self, pose):
    method _matrix_to_pose (line 398) | def _matrix_to_pose(self, matrix, offset=np.eye(4)):
    method _pose_to_matrix (line 416) | def _pose_to_matrix(self, pose):
    method _get_low_dim_data (line 431) | def _get_low_dim_data(self, obs):
    method __len__ (line 447) | def __len__(self):
    method min_traj_len (line 450) | def min_traj_len(self):
    method get_conditions (line 454) | def get_conditions(self, observations):
    method _calc_rank (line 464) | def _calc_rank(self, poses):
    method _filter_traj (line 475) | def _filter_traj(self, joints_positions: np.ndarray) -> list:
    method __getitem__ (line 496) | def __getitem__(self, index) -> Any:
    method _prepare_input (line 565) | def _prepare_input(self, task: str, use_demos: bool = False):
    method _create_uniform_pixel_coords_image (line 623) | def _create_uniform_pixel_coords_image(self, resolution: np.ndarray):
    method _transform (line 647) | def _transform(self, coords, trans):
    method _pixel_to_world_coords (line 668) | def _pixel_to_world_coords(self, pixel_coords, cam_proj_mat_inv):
    method _pointcloud_from_depth_and_camera_params (line 685) | def _pointcloud_from_depth_and_camera_params(

FILE: rk_diffuser/dataset/rl_bench_dataset.py
  class NormEEPosePlanning (line 35) | class NormEEPosePlanning(EndEffectorPoseViaPlanning):
    method action (line 36) | def action(self, scene: Scene, action: np.ndarray):
  class TrajectoryActionMode (line 46) | class TrajectoryActionMode(EndEffectorPoseViaIK):
    method __init__ (line 49) | def __init__(
    method action (line 59) | def action(self, scene: Scene, action: np.ndarray):
    method action_shape (line 71) | def action_shape(self, scene: Scene) -> tuple:
  class JointsTrajectoryActionMode (line 75) | class JointsTrajectoryActionMode(ArmActionMode):
    method __init__ (line 78) | def __init__(self, points: int):
    method action (line 81) | def action(self, scene: Scene, action: np.ndarray):
    method _pre_proc_traj (line 97) | def _pre_proc_traj(self, action):
    method action_shape (line 109) | def action_shape(self, scene: Scene) -> tuple:
  function _create_obs_config (line 113) | def _create_obs_config(camera_names: List[str], camera_resolution: List[...
  function _get_action_mode (line 152) | def _get_action_mode(num_points, use_traj=False):
  class RLBenchDataset (line 161) | class RLBenchDataset(Dataset):
    method __init__ (line 162) | def __init__(
    method _get_env (line 263) | def _get_env(self, task, headless=True, use_traj=False):
    method _get_demos (line 291) | def _get_demos(self, task, ret_dict):
    method _load_demos_to_dataset (line 317) | def _load_demos_to_dataset(self, demos, num_of_episodes, ds_list):
    method __len__ (line 399) | def __len__(self):
    method min_traj_len (line 402) | def min_traj_len(self):
    method get_conditions (line 406) | def get_conditions(self, observations):
    method _calc_rank (line 416) | def _calc_rank(self, poses):
    method __getitem__ (line 428) | def __getitem__(self, index) -> Any:

FILE: rk_diffuser/dataset/rl_bench_env.py
  class CustomRLBenchEnv (line 20) | class CustomRLBenchEnv(RLBenchEnv):
    method __init__ (line 21) | def __init__(
    method observation_elements (line 52) | def observation_elements(self) -> List[ObservationElement]:
    method extract_obs (line 63) | def extract_obs(self, obs: Observation, t=None, prev_action=None):
    method launch (line 94) | def launch(self):
    method reset (line 105) | def reset(self) -> dict:
    method register_callback (line 113) | def register_callback(self, func):
    method _my_callback (line 116) | def _my_callback(self):
    method _append_final_frame (line 121) | def _append_final_frame(self, success: bool):
    method step (line 130) | def step(self, action: np.ndarray, record: bool = False) -> Transition:
    method reset_to_demo (line 166) | def reset_to_demo(self, demo):
    method eval_planner (line 172) | def eval_planner(self, amount, callable_each_step):
    method _get_live_demos (line 180) | def _get_live_demos(

FILE: rk_diffuser/models/diffusion.py
  function _create_model_fn (line 20) | def _create_model_fn(
  class GaussianDynDiffusion (line 62) | class GaussianDynDiffusion(nn.Module):
    method __init__ (line 63) | def __init__(
    method gripper_pose_loss (line 252) | def gripper_pose_loss(self, p1, p2):
    method get_loss_weights (line 271) | def get_loss_weights(self, discount: float) -> torch.Tensor:
    method predict_start_from_noise (line 296) | def predict_start_from_noise(
    method q_posterior (line 318) | def q_posterior(
    method p_mean_variance (line 343) | def p_mean_variance(
    method p_sample (line 414) | def p_sample(
    method p_sample_loop (line 441) | def p_sample_loop(
    method conditional_sample (line 512) | def conditional_sample(
    method q_sample (line 538) | def q_sample(
    method init_noise (line 562) | def init_noise(
    method p_losses (line 607) | def p_losses(
    method proc_cond (line 718) | def proc_cond(self, cond: dict) -> dict:
    method loss (line 735) | def loss(self, x: torch.tensor, cond: dict, **kwargs) -> tuple:
    method forward (line 753) | def forward(self, cond: dict, *args, **kwargs) -> dict:

FILE: rk_diffuser/models/helpers.py
  class SinusoidalPosEmb (line 18) | class SinusoidalPosEmb(nn.Module):
    method __init__ (line 19) | def __init__(self, dim):
    method forward (line 23) | def forward(self, x):
  class Downsample1d (line 33) | class Downsample1d(nn.Module):
    method __init__ (line 34) | def __init__(self, dim):
    method forward (line 38) | def forward(self, x):
  class Upsample1d (line 42) | class Upsample1d(nn.Module):
    method __init__ (line 43) | def __init__(self, dim):
    method forward (line 47) | def forward(self, x):
  class Conv1dBlock (line 51) | class Conv1dBlock(nn.Module):
    method __init__ (line 56) | def __init__(self, inp_channels, out_channels, kernel_size, mish=True,...
    method forward (line 74) | def forward(self, x):
  function extract (line 83) | def extract(a, t, x_shape):
  function cosine_beta_schedule (line 89) | def cosine_beta_schedule(timesteps, s=0.008, dtype=torch.float32):
  function apply_conditioning (line 103) | def apply_conditioning(x, conditions, action_dim):
  class WeightedLoss (line 114) | class WeightedLoss(nn.Module):
    method __init__ (line 115) | def __init__(self, weights, action_dim):
    method forward (line 120) | def forward(self, pred, targ):
  class WeightedStateLoss (line 133) | class WeightedStateLoss(nn.Module):
    method __init__ (line 134) | def __init__(self, weights):
    method forward (line 138) | def forward(self, pred, targ):
  class ValueLoss (line 148) | class ValueLoss(nn.Module):
    method __init__ (line 149) | def __init__(self, *args):
    method forward (line 153) | def forward(self, pred, targ):
  class WeightedL1 (line 176) | class WeightedL1(WeightedLoss):
    method _loss (line 177) | def _loss(self, pred, targ):
  class WeightedL2 (line 181) | class WeightedL2(WeightedLoss):
    method _loss (line 182) | def _loss(self, pred, targ):
  class WeightedStateL2 (line 186) | class WeightedStateL2(WeightedStateLoss):
    method _loss (line 187) | def _loss(self, pred, targ):
  class WeightedChamfer (line 191) | class WeightedChamfer(WeightedStateLoss):
    method __init__ (line 192) | def __init__(self, weights, converage_weight=1.0, norm="l1"):
    method _loss (line 197) | def _loss(self, pred, targ):
  class ValueL1 (line 212) | class ValueL1(ValueLoss):
    method _loss (line 213) | def _loss(self, pred, targ):
  class ValueL2 (line 217) | class ValueL2(ValueLoss):
    method _loss (line 218) | def _loss(self, pred, targ):

FILE: rk_diffuser/models/multi_level_diffusion.py
  class MultiLevelDiffusion (line 10) | class MultiLevelDiffusion(nn.Module):
    method __init__ (line 18) | def __init__(
    method robot_offset (line 50) | def robot_offset(self):
    method robot_rot (line 55) | def robot_rot(self):
    method conditional_sample (line 61) | def conditional_sample(
    method loss (line 120) | def loss(self, x: torch.tensor, cond: dict, **kwargs) -> tuple:
    method _form_diffusion_batch (line 144) | def _form_diffusion_batch(

FILE: rk_diffuser/models/pointnet.py
  class STN3d (line 12) | class STN3d(nn.Module):
    method __init__ (line 13) | def __init__(self):
    method forward (line 32) | def forward(self, x):
  class STNkd (line 60) | class STNkd(nn.Module):
    method __init__ (line 61) | def __init__(self, k=64):
    method forward (line 79) | def forward(self, x):
  class PointNetfeat (line 103) | class PointNetfeat(nn.Module):
    method __init__ (line 104) | def __init__(self, feature_transform=False):
    method forward (line 117) | def forward(self, x):
  class PointNetCls (line 143) | class PointNetCls(nn.Module):
    method __init__ (line 144) | def __init__(self, k=2, feature_transform=False):
    method forward (line 156) | def forward(self, x):
  class PointNetDenseCls (line 164) | class PointNetDenseCls(nn.Module):
    method __init__ (line 165) | def __init__(self, k=2, feature_transform=False):
    method forward (line 178) | def forward(self, x):

FILE: rk_diffuser/models/resnet.py
  function preproc_fn (line 13) | def preproc_fn(h=224, w=224, rgb=True):
  function infer_model_out_shape (line 38) | def infer_model_out_shape(model, input_shape: Tuple) -> tuple:
  class ResnetEncoder (line 51) | class ResnetEncoder(nn.Module):
    method __init__ (line 54) | def __init__(
    method forward (line 90) | def forward(self, img: torch.Tensor) -> torch.Tensor:

FILE: rk_diffuser/models/temporal.py
  class Residual (line 20) | class Residual(nn.Module):
    method __init__ (line 21) | def __init__(self, fn):
    method forward (line 25) | def forward(self, x, *args, **kwargs):
  class PreNorm (line 29) | class PreNorm(nn.Module):
    method __init__ (line 30) | def __init__(self, dim, fn):
    method forward (line 35) | def forward(self, x):
  class LinearAttention (line 40) | class LinearAttention(nn.Module):
    method __init__ (line 41) | def __init__(self, dim, heads=4, dim_head=128):
    method forward (line 48) | def forward(self, x):
  class GlobalMixing (line 63) | class GlobalMixing(nn.Module):
    method __init__ (line 64) | def __init__(self, dim, heads=4, dim_head=128):
    method forward (line 71) | def forward(self, x):
  class ResidualTemporalBlock (line 86) | class ResidualTemporalBlock(nn.Module):
    method __init__ (line 87) | def __init__(
    method forward (line 116) | def forward(self, x, t):
  class Temporal (line 129) | class Temporal(nn.Module):
    method __init__ (line 130) | def __init__(
    method forward (line 279) | def forward(
  class UnetBackbone (line 326) | class UnetBackbone(nn.Module):
    method __init__ (line 329) | def __init__(
    method forward (line 451) | def forward(self, x: torch.tensor, t: torch.tensor) -> torch.tensor:
  class TransformerBackbone (line 480) | class TransformerBackbone(nn.Module):
    method __init__ (line 481) | def __init__(
    method forward (line 567) | def forward(self, x: torch.tensor, t: torch.tensor) -> torch.tensor:

FILE: rk_diffuser/robot.py
  class DiffRobot (line 11) | class DiffRobot(nn.Module):
    method __init__ (line 14) | def __init__(
    method to (line 59) | def to(self, device: torch.device) -> None:
    method forward_kinematics_batch (line 72) | def forward_kinematics_batch(self, joints: torch.tensor) -> torch.tensor:
    method inverse_kinematics_autodiff_single_step_batch_pt (line 90) | def inverse_kinematics_autodiff_single_step_batch_pt(

FILE: rk_diffuser/trainer.py
  function cycle (line 36) | def cycle(dl):
  class EMA (line 42) | class EMA:
    method __init__ (line 47) | def __init__(self, beta):
    method update_model_average (line 51) | def update_model_average(self, ma_model, current_model):
    method _update_average (line 58) | def _update_average(self, old, new):
  class Trainer (line 64) | class Trainer(object):
    method __init__ (line 65) | def __init__(
    method create_envs (line 181) | def create_envs(self) -> None:
    method _set_logger (line 255) | def _set_logger(self) -> None:
    method _reset_parameters (line 285) | def _reset_parameters(self) -> None:
    method _step_ema (line 291) | def _step_ema(self):
    method generate_traj (line 310) | def generate_traj(self, **kwargs) -> dict:
    method get_eval_log (line 333) | def get_eval_log(
    method train (line 434) | def train(self, n_train_steps: int, train: bool = True) -> None:
    method _prepare_demo_data (line 511) | def _prepare_demo_data(self, task: str) -> tuple:
    method _dispatch_eval_trajs (line 608) | def _dispatch_eval_trajs(self, trajs: list, task_names: list) -> dict:
    method online_eval (line 626) | def online_eval(self) -> dict:
    method __exit__ (line 714) | def __exit__():

FILE: rk_diffuser/utils.py
  class Timer (line 18) | class Timer:
    method __init__ (line 19) | def __init__(self):
    method __call__ (line 22) | def __call__(self, reset=True):
  class Progress (line 30) | class Progress:
    method __init__ (line 31) | def __init__(
    method update (line 62) | def update(self, description, n=1):
    method resume (line 69) | def resume(self):
    method pause (line 75) | def pause(self):
    method set_description (line 79) | def set_description(self, params=[]):
    method append_description (line 112) | def append_description(self, descr):
    method _clear (line 115) | def _clear(self):
    method _format_percent (line 122) | def _format_percent(self, n, total):
    method _format_speed (line 140) | def _format_speed(self, n):
    method _chunk (line 149) | def _chunk(self, l, n):
    method _format (line 152) | def _format(self, chunks):
    method _format_chunk (line 159) | def _format_chunk(self, chunk):
    method _format_param (line 163) | def _format_param(self, param):
    method stamp (line 167) | def stamp(self):
    method close (line 180) | def close(self):
  class Silent (line 184) | class Silent:
    method __init__ (line 185) | def __init__(self, *args, **kwargs):
    method __getattr__ (line 188) | def __getattr__(self, attr):
  function to_np (line 192) | def to_np(x):
  function to_torch (line 198) | def to_torch(x, dtype=None, device=None):
  function to_device (line 209) | def to_device(x, device=DEVICE):
  function batchify (line 232) | def batchify(batch, device):
  function apply_dict (line 248) | def apply_dict(fn, d, *args, **kwargs):
  function normalize (line 252) | def normalize(x):
  function to_img (line 261) | def to_img(x):
  function set_device (line 268) | def set_device(device):
  function batch_to_device (line 274) | def batch_to_device(batch, device="cuda:0"):
  function _to_str (line 278) | def _to_str(num):
  function param_to_module (line 290) | def param_to_module(param):
  function report_parameters (line 295) | def report_parameters(model, topk=10):
  function _is_stopped (line 316) | def _is_stopped(demo, i, obs, stopped_buffer, delta=0.1):
  function keypoint_discovery (line 333) | def keypoint_discovery(demo: Demo, stopping_delta=0.1) -> List[int]:
  function cumsum_traj (line 354) | def cumsum_traj(x, cond):
  function concatenate_tensors_dict_list (line 362) | def concatenate_tensors_dict_list(dict_list, dim=0):
  function geodesic_distance_between_quaternions (line 382) | def geodesic_distance_between_quaternions(
  function diff_kinematics_pose_to_coppelia (line 418) | def diff_kinematics_pose_to_coppelia(diff_k_pose, offset):
  function proc_quaternion (line 435) | def proc_quaternion(poses):
  function _sqrt_positive_part (line 446) | def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
  function matrix_to_quaternion (line 457) | def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
  function rlb_pose_to_matrix (line 519) | def rlb_pose_to_matrix(pose):
  function matrix_to_rlb_pose (line 531) | def matrix_to_rlb_pose(matrix):
  function load_low_level_ckpt (line 540) | def load_low_level_ckpt(path):
  function load_checkpoint (line 549) | def load_checkpoint(diffusion, load_model_path, backbone):

FILE: train_low_level.py
  function _create_agent_fn (line 17) | def _create_agent_fn(
  function main (line 73) | def main(cfgs):