File size: 2,066 Bytes
dc9f917 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """Verify the frozen LeWM loads, encodes, and gives gradients w.r.t. actions."""
import sys
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from lejepa_control.world_model import load_lewm # noqa: E402
def main():
import stable_worldmodel as swm
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = load_lewm(device=device)
print('loaded LeWM ok')
print(' predictor.num_frames =', model.predictor.num_frames)
print(' action_encoder.input_dim =', model.action_encoder.input_dim)
ds = swm.data.load_dataset('pusht_smoke')
print(' dataset columns =', ds.column_names)
print(' episodes =', len(ds.lengths))
ep = ds.load_episode(0)
for k, v in ep.items():
if hasattr(v, 'shape'):
print(f' ep[{k}] {tuple(v.shape)} {v.dtype}')
pixels_key = 'pixels' if 'pixels' in ep else 'obs.pixels'
frames = ep[pixels_key]
if not torch.is_tensor(frames):
frames = torch.as_tensor(frames)
if frames.shape[-1] in (1, 3): # NHWC -> NCHW
frames = frames.permute(0, 3, 1, 2)
frames = frames.float() / 255.0 if frames.dtype == torch.uint8 else frames
T = model.predictor.num_frames
pixels = frames[:T].unsqueeze(0).to(device) # (1, T, C, H, W)
with torch.no_grad():
info = model.encode({'pixels': pixels})
emb = info['emb']
print(' emb', tuple(emb.shape))
# gradient check: does d(pred)/d(action) flow?
action_dim = model.action_encoder.input_dim
action = torch.zeros(1, T, action_dim, device=device, requires_grad=True)
act_emb = model.action_encoder(action)
pred = model.predict(emb, act_emb)
print(' pred', tuple(pred.shape))
pred.sum().backward()
g = action.grad
print(' action.grad norm =', float(g.norm()))
assert g.norm() > 0, 'no gradient reached the action'
print('OK: gradients flow through the frozen predictor to actions')
if __name__ == '__main__':
main()
|