stable_worldmodel.wm package

Submodules

stable_worldmodel.wm.dinowm module

class Attention(dim, heads=8, dim_head=64, dropout=0.0, num_patches=1, num_frames=1)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

generate_mask_matrix(npatch, nwindow)

class CausalPredictor(*, num_patches, num_frames, dim, depth, heads, mlp_dim, pool='cls', dim_head=64, dropout=0.0, emb_dropout=0.0)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class DINOWM(encoder, predictor, action_encoder, proprio_encoder, decoder=None, history_size=3, num_pred=1, device='cpu')

Bases: Module

decode(info)

encode(info, pixels_key='pixels', target='embed', proprio_key=None, action_key=None)

get_cost(info_dict: dict, action_candidates: Tensor)

predict(embedding)

predict next latent state :param embedding: (B, T, P, d)

Returns:

(B, T, P, d)

Return type:

preds

replace_action_in_embedding(embedding, act)

Replace the action embeddings in the latent state z with the provided actions.

rollout(info, action_sequence)

Rollout the world model given an initial observation and a sequence of actions.

Params: obs_start: n current observations (B, n, C, H, W) actions: current and predicted actions (B, n+t, action_dim)

Returns: z_obs: dict with latent observations (B, n+t+1, n_patches, D) z: predicted latent states (B, n+t+1, n_patches, D)

split_embedding(embedding, action_dim, proprio_dim)

class Decoder(in_channel, out_channel, channel, n_res_block, n_res_channel, stride)

Bases: Module

forward(input)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class Embedder(num_frames=1, tubelet_size=1, in_chans=8, emb_dim=10)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class Encoder(in_channel, channel, n_res_block, n_res_channel, stride)

Bases: Module

forward(input)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class FeedForward(dim, hidden_dim, dropout=0.0)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class Quantize(dim, n_embed, decay=0.99, eps=1e-05)

Bases: Module

embed_code(embed_id)

forward(input)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ResBlock(in_channel, channel)

Bases: Module

forward(input)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class Transformer(dim, depth, heads, dim_head, mlp_dim, dropout=0.0, num_patches=1, num_frames=1)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class VQVAE(in_channel=3, channel=128, n_res_block=2, n_res_channel=32, emb_dim=64, n_embed=512, decay=0.99, quantize=True)

Bases: Module

decode(quant_b)

decode_code(code_b)

forward(input)

input: (b, t, num_patches, emb_dim)

all_reduce(tensor, op=<RedOpType.SUM: 0>)

get_world_size()

stable_worldmodel.wm.dreamer module

stable_worldmodel.wm.dummy module

class DummyWorldModel(image_shape, action_dim)

Bases: Module

encode(obs)

get_cost(info_dict: dict, action_candidates: Tensor)

predict(obs, actions, timestep=None)

Predict next s_t+H embedding given s_t + action sequence i.e rollout the dynamics model for H steps.

transform(info_dict)

stable_worldmodel.wm.frame module

stable_worldmodel.wm.tdmpc module

Module contents

class DINOWM(encoder, predictor, action_encoder, proprio_encoder, decoder=None, history_size=3, num_pred=1, device='cpu')

Bases: Module

decode(info)

encode(info, pixels_key='pixels', target='embed', proprio_key=None, action_key=None)

get_cost(info_dict: dict, action_candidates: Tensor)

predict(embedding)

predict next latent state :param embedding: (B, T, P, d)

Returns:

(B, T, P, d)

Return type:

preds

replace_action_in_embedding(embedding, act)

Replace the action embeddings in the latent state z with the provided actions.

rollout(info, action_sequence)

Rollout the world model given an initial observation and a sequence of actions.

Params: obs_start: n current observations (B, n, C, H, W) actions: current and predicted actions (B, n+t, action_dim)

Returns: z_obs: dict with latent observations (B, n+t+1, n_patches, D) z: predicted latent states (B, n+t+1, n_patches, D)

split_embedding(embedding, action_dim, proprio_dim)

class DummyWorldModel(image_shape, action_dim)

Bases: Module

encode(obs)

get_cost(info_dict: dict, action_candidates: Tensor)

predict(obs, actions, timestep=None)

Predict next s_t+H embedding given s_t + action sequence i.e rollout the dynamics model for H steps.