API reference

Generated from the docstrings of torch_world.

module	contents
torch_world.nn	ready encoders CLSEncoder, PixelPatches and PatchTokens (DINOv2), TextEncoder and WordTokenizer; VisionTransformer (image and video), MaskedPredictor, Transformer, Block, AdaLNBlock, Attention, MLP, sin-cos and rotary positions
torch_world.masks	MultiBlockMask (2D), MultiBlock3DMask (video), frame_causal_mask, apply_masks
torch_world.losses	SIGReg, info_nce
torch_world.models	the JEPA and WorldModel base classes and the models
torch_world.planning	CEM, for any WorldModel, scoring the final or every imagined step
torch_world.envs	RobotCube (reach or pick-and-place) and PushT (MuJoCo), the pixel worlds Macaw and PointMass; run_policy
torch_world.datasets	TrajectoryDataset, MedMNIST, MedicalVQA
torch_world.evaluation	protocols that run a model and score it: goal_reaching for planners (success rates, videos), vqa and compare for question answering
torch_world.metrics	the metrics of the papers, on plain tensors: linear, k-NN and attentive probes; RankMe and LiDAR; rollout error, SSIM, PSNR and Chamfer distance; recall@k and top-k accuracy; for robotics, success rates with confidence intervals, steps to success, progress, SPL and smoothness
torch_world.utils	ema_update, linear_schedule, recompute_batchnorm
torch_world.plots	matplotlib figures for notebooks (optional dependency)