Add post_pos as a model target via travel_dir decomposition

step_length already encodes |post_pos - pre_pos| by definition, so a raw
post_pos target would duplicate that magnitude and could drift inconsistent
with step_length during sampling. Instead add travel_dir, a unit vector
(local frame) giving only the direction of pre_pos->post_pos; post_pos is
reconstructed at inference as pre_pos + step_length * travel_dir, keeping
the two self-consistent. Target grows from 6D to 9D.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-06-18 10:36:55 +02:00
parent c3b7b2744c
commit 72bd65ff9f
11 changed files with 121 additions and 17 deletions
+3 -3
View File
@@ -18,7 +18,7 @@ def test_sinusoidal_embedding_batch_1():
def test_denoising_mlp_output_shape():
B = 8
model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
x_t = torch.randn(B, 6)
x_t = torch.randn(B, 9)
t = torch.rand(B)
cond_cont = torch.randn(B, 9)
cond_cat = torch.stack([
@@ -26,13 +26,13 @@ def test_denoising_mlp_output_shape():
torch.randint(0, 3, (B,)),
], dim=1)
out = model(x_t, t, cond_cont, cond_cat)
assert out.shape == (B, 6)
assert out.shape == (B, 9)
def test_denoising_mlp_gradients_flow():
B = 4
model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
x_t = torch.randn(B, 6)
x_t = torch.randn(B, 9)
t = torch.rand(B)
cond_cont = torch.randn(B, 9)
cond_cat = torch.zeros(B, 2, dtype=torch.long)