| from __future__ import annotations |
|
|
| import importlib |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| from experiments.shared.src.methods import PAPER_LEARNED_METHODS |
| from experiments.shared.src.vision.clean_renderer import render_clean_boat_array |
|
|
|
|
| def make_batch( |
| batch_size: int = 2, |
| history_len: int = 32, |
| horizon: int = 4, |
| image_size: int = 160, |
| action_dim: int = 3, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| rng = np.random.default_rng(11) |
| images = np.zeros((batch_size, history_len, 3, image_size, image_size), dtype=np.uint8) |
| actions = rng.uniform(-0.3, 0.3, size=(batch_size, history_len, action_dim)).astype(np.float32) |
| future_actions = rng.uniform(-0.3, 0.3, size=(batch_size, horizon, action_dim)).astype(np.float32) |
| targets = np.zeros((batch_size, horizon, 4), dtype=np.float32) |
| for b in range(batch_size): |
| x = float(rng.uniform(2.0, 8.0)) |
| y = float(rng.uniform(2.0, 8.0)) |
| theta = float(rng.uniform(-np.pi, np.pi)) |
| for t in range(history_len): |
| state = np.array([x + 0.01 * t, y, theta + 0.005 * t, 0.0, 0.0, 0.0], dtype=np.float32) |
| image = render_clean_boat_array(state, "twin", image_size=image_size, visual_scale=2.5) |
| images[b, t] = np.transpose(image, (2, 0, 1)) |
| for h in range(horizon): |
| theta_h = theta + 0.005 * (history_len + h) |
| targets[b, h] = np.array( |
| [x + 0.01 * (history_len + h), y, np.cos(theta_h), np.sin(theta_h)], |
| dtype=np.float32, |
| ) |
| return ( |
| torch.from_numpy(images), |
| torch.from_numpy(actions), |
| torch.from_numpy(future_actions), |
| torch.from_numpy(targets), |
| ) |
|
|
|
|
| def test_all_paper_image_world_models_train_step() -> None: |
| batch = make_batch() |
| images, actions, future_actions, targets = batch |
| for method in PAPER_LEARNED_METHODS: |
| cfg = importlib.import_module(f"experiments.{method}.src.config").default_config() |
| model = importlib.import_module(f"experiments.{method}.src.model").build_model(cfg) |
| pred = model.rollout(images, actions, future_actions) |
| assert pred.shape == targets.shape |
| loss = F.mse_loss(pred, targets) |
| loss.backward() |
| assert all(p.grad is not None for p in model.parameters() if p.requires_grad) |
|
|
|
|
| def test_flowmo_context_residual_is_action_separated() -> None: |
| cfg = importlib.import_module("experiments.flowmo.src.config").default_config() |
| model = importlib.import_module("experiments.flowmo.src.model").build_model(cfg) |
| z = torch.randn(4, cfg.z_dim) |
| c = torch.randn(4, cfg.c_dim) |
| a1 = torch.randn(4, cfg.action_dim).clamp(-1.0, 1.0) |
| a2 = torch.randn(4, cfg.action_dim).clamp(-1.0, 1.0) |
| z1 = model.step(z, a1, c) - model.step(z, a1, torch.zeros_like(c)) |
| z2 = model.step(z, a2, c) - model.step(z, a2, torch.zeros_like(c)) |
| assert torch.allclose(z1, z2, atol=1e-6) |
|
|
