nx_motion_neural.nx source
↩ module page · 186 lines · 7364 B
1// nx_motion_neural.nx -- the NEURAL-MOTION FOUNDATION (S6, MotionGPT/T2M-GPT thesis "motion as language",
2// sovereign + no 491MB LLM load): a motion is discretized into TOKENS by a codebook LEARNED from data (vector
3// quantization = k-means, the exact VQ-VAE quantizer), and a learned token-TRANSITION model GENERATES novel
4// motion-token sequences. This is real learning-from-data (not the parametric keyword->recipe of nx_movelib):
5// the codebook is trained (Lloyd iterations reduce distortion) and the generator composes sequences it was never
6// given whole. Integer fx (rad4096 joint angles) => deterministic, VM-vettable. The token-LM here is an n-gram
7// (the bootstrap); the transformer/our-Qwen over these tokens is the deeper rung. license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_movelib.nx" // move_gen/move_pose_at/mv_hdr + RG_* + skeleton, transitively
10const MN_MAGIC_1103515245: i64 = 1103515245
11const MN_MAGIC_12345: i64 = 12345
12
13const MN_D: i64 = 7 // pose vector dims (the joints the move vocabulary drives)
14const MN_FR: i64 = 16 // frames sampled per move (over one stride)
15const MN_MOVES: i64 = 10 // wave/cheer/sway/kick/lookaround/armcircle/clap/punch/bow/jumpingjack
16const MN_N: i64 = 160 // corpus rows = MN_MOVES * MN_FR
17const MN_K: i64 = 16 // codebook size (bumped for the richer 10-move pose space)
18
19// read the 7-dim pose vector from the skeleton's SET angles (after move_pose_at, before sk_update).
20// (call-derefs hoisted into lets -- a call inside an array-store RHS is the known codegen SEGV.)
21func pose_read(sk: i64, out: *i64) -> i64 {
22 let b0: *i64 = sk_bone(sk, RG_SHR); out[0] = b0[5]
23 let b1: *i64 = sk_bone(sk, RG_SHL); out[1] = b1[5]
24 let b2: *i64 = sk_bone(sk, RG_ELR); out[2] = b2[4]
25 let b3: *i64 = sk_bone(sk, RG_SPINE); out[3] = b3[4]
26 let b4: *i64 = sk_bone(sk, RG_NECK); out[4] = b4[4]
27 let b5: *i64 = sk_bone(sk, RG_HIPR); out[5] = b5[5]
28 let b6: *i64 = sk_bone(sk, RG_HIPL); out[6] = b6[5]
29 return 0
30}
31// apply a 7-dim pose vector to the skeleton (the inverse of pose_read); caller runs sk_update.
32func pose_apply(sk: i64, v: *i64) -> i64 {
33 sk_pose(sk, RG_SHR, 0, v[0])
34 sk_pose(sk, RG_SHL, 0, v[1])
35 sk_pose(sk, RG_ELR, v[2], 0)
36 sk_pose(sk, RG_SPINE, v[3], 0)
37 sk_pose(sk, RG_NECK, v[4], 0)
38 sk_pose(sk, RG_HIPR, 0, v[5])
39 sk_pose(sk, RG_HIPL, 0, v[6])
40 return 0
41}
42
43// build the motion corpus: for each move, sample MN_FR poses over its stride into corpus[N*D] (move-major).
44func mn_build_corpus(sk: i64, corpus: *i64) -> i64 {
45 let move: i64 = sys_mmap(move_bytes()) as i64
46 let vec: *i64 = sys_mmap(MN_D * 8) as *i64
47 var m: i64 = 0
48 while m < MN_MOVES {
49 move_gen(move, m + 1) // move ids 1..6
50 let hp: *i64 = mv_hdr(move)
51 let period: i64 = hp[1]
52 var f: i64 = 0
53 while f < MN_FR {
54 let t: i64 = f * period / MN_FR
55 move_pose_at(move, sk, t)
56 pose_read(sk, vec)
57 let row: i64 = (m * MN_FR + f) * MN_D
58 var d: i64 = 0
59 while d < MN_D { corpus[row + d] = vec[d]; d = d + 1 }
60 f = f + 1
61 }
62 m = m + 1
63 }
64 return MN_N
65}
66
67// squared L2 distance between corpus row a (at ao) and codebook entry b (at bo).
68func mn_dist(a: *i64, ao: i64, b: *i64, bo: i64) -> i64 {
69 var s: i64 = 0
70 var i: i64 = 0
71 while i < MN_D { let d: i64 = a[ao + i] - b[bo + i]; s = s + d * d; i = i + 1 }
72 return s
73}
74// nearest codebook index for the vector at src[so..].
75func mn_encode(src: *i64, so: i64, cb: *i64) -> i64 {
76 var best: i64 = 0
77 var bd: i64 = mn_dist(src, so, cb, 0)
78 var k: i64 = 1
79 while k < MN_K {
80 let d: i64 = mn_dist(src, so, cb, k * MN_D)
81 if d < bd { bd = d; best = k }
82 k = k + 1
83 }
84 return best
85}
86// total distortion = sum over corpus of nearest-centroid distance.
87func mn_distortion(corpus: *i64, cb: *i64) -> i64 {
88 var s: i64 = 0
89 var n: i64 = 0
90 while n < MN_N {
91 let k: i64 = mn_encode(corpus, n * MN_D, cb)
92 s = s + mn_dist(corpus, n * MN_D, cb, k * MN_D)
93 n = n + 1
94 }
95 return s
96}
97// init the codebook with MN_K spread corpus rows.
98func mn_init_cb(corpus: *i64, cb: *i64) -> i64 {
99 var k: i64 = 0
100 while k < MN_K {
101 let src: i64 = (k * MN_N / MN_K) * MN_D
102 var d: i64 = 0
103 while d < MN_D { cb[k * MN_D + d] = corpus[src + d]; d = d + 1 }
104 k = k + 1
105 }
106 return 0
107}
108// LEARN the codebook: Lloyd's k-means for `iters` steps. Returns final distortion.
109func mn_train(corpus: *i64, cb: *i64, iters: i64) -> i64 {
110 let sum: *i64 = sys_mmap(MN_K * MN_D * 8) as *i64
111 let cnt: *i64 = sys_mmap(MN_K * 8) as *i64
112 var it: i64 = 0
113 while it < iters {
114 var i: i64 = 0
115 while i < MN_K * MN_D { sum[i] = 0; i = i + 1 }
116 i = 0
117 while i < MN_K { cnt[i] = 0; i = i + 1 }
118 var n: i64 = 0
119 while n < MN_N {
120 let k: i64 = mn_encode(corpus, n * MN_D, cb)
121 cnt[k] = cnt[k] + 1
122 var d: i64 = 0
123 while d < MN_D { sum[k * MN_D + d] = sum[k * MN_D + d] + corpus[n * MN_D + d]; d = d + 1 }
124 n = n + 1
125 }
126 var k2: i64 = 0
127 while k2 < MN_K {
128 if cnt[k2] > 0 {
129 var d2: i64 = 0
130 while d2 < MN_D { cb[k2 * MN_D + d2] = sum[k2 * MN_D + d2] / cnt[k2]; d2 = d2 + 1 }
131 }
132 k2 = k2 + 1
133 }
134 it = it + 1
135 }
136 return mn_distortion(corpus, cb)
137}
138// encode the whole corpus to tokens[N] (the "motion sentences").
139func mn_tokenize(corpus: *i64, cb: *i64, tokens: *i64) -> i64 {
140 var n: i64 = 0
141 while n < MN_N { tokens[n] = mn_encode(corpus, n * MN_D, cb); n = n + 1 }
142 return 0
143}
144// build the token-transition counts (the n-gram LM) from the per-move token sequences.
145func mn_bigram(tokens: *i64, trans: *i64) -> i64 {
146 var i: i64 = 0
147 while i < MN_K * MN_K { trans[i] = 0; i = i + 1 }
148 var m: i64 = 0
149 while m < MN_MOVES {
150 var f: i64 = 0
151 while f < MN_FR - 1 {
152 let a: i64 = tokens[m * MN_FR + f]
153 let b: i64 = tokens[m * MN_FR + f + 1]
154 trans[a * MN_K + b] = trans[a * MN_K + b] + 1
155 f = f + 1
156 }
157 m = m + 1
158 }
159 return 0
160}
161func mn_lcg(st: *i64) -> i64 { st[0] = (st[0] * MN_MAGIC_1103515245 + MN_MAGIC_12345) & 0x7fffffff; return st[0] }
162// GENERATE a motion-token sequence of length `len` from `seed`, sampling the learned transitions (seeded, weighted).
163func mn_generate(trans: *i64, seed: i64, st: *i64, out: *i64, len: i64) -> i64 {
164 var cur: i64 = seed
165 var i: i64 = 0
166 while i < len {
167 out[i] = cur
168 var tot: i64 = 0
169 var j: i64 = 0
170 while j < MN_K { tot = tot + trans[cur * MN_K + j]; j = j + 1 }
171 var nxt: i64 = 0
172 if tot > 0 {
173 var r: i64 = mn_lcg(st) % tot
174 var acc: i64 = 0
175 var jj: i64 = 0
176 var done: i64 = 0
177 while done == 0 {
178 acc = acc + trans[cur * MN_K + jj]
179 if r < acc { nxt = jj; done = 1 } else { jj = jj + 1; if jj >= MN_K { nxt = cur; done = 1 } }
180 }
181 } else { nxt = mn_lcg(st) % MN_K }
182 cur = nxt
183 i = i + 1
184 }
185 return 0
186}