mjswan 0.8.1 → 0.8.2

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "mjswan",
   "type": "module",
-  "version": "0.8.1",
+  "version": "0.8.2",
   "repository": {
     "type": "git",
     "url": "https://github.com/ttktjmt/mjswan.git"

package/src/core/dsl/__tests__/primitives.test.ts

@@ -0,0 +1,98 @@
+/**
+ * Direct unit tests for the DSL primitives added for composition-graph
+ * observations (ADR 0003): arithmetic (Div/Sqrt/Sum), Slice, and the
+ * motion-reference sources (TrackingRefField/TrackingIsReady).  Each primitive
+ * is exercised on its own, not through any task's observation graph.
+ */
+import { describe, expect, it } from 'vitest';
+
+import type { EvalContext } from '../interpreter';
+import { Primitives } from '../primitives';
+
+function expectClose(a: Float32Array, b: ArrayLike<number>, tol = 1e-6): void {
+  expect(a.length).toBe(b.length);
+  for (let i = 0; i < a.length; i++) expect(Math.abs(a[i] - b[i])).toBeLessThan(tol);
+}
+
+// A minimal motion command exposing the reference-trajectory buffers the
+// sources read (the generic MotionRefSource shape — not tied to any task).
+interface RefTerm {
+  isReady(): boolean;
+  refRootPos: Float32Array[];
+  refRootQuat: Float32Array[];
+  refJointPos: Float32Array[];
+  refIdx: number;
+  refLen: number;
+}
+
+function ctx(term: RefTerm | null, nJoints = 3): EvalContext {
+  const runner = {
+    getContext: () => ({
+      commandManager: { getTerm: (n: string) => (n === 'motion' ? term : null) },
+    }),
+    getNumActions: () => nJoints,
+    getLastActions: () => new Float32Array(nJoints),
+  };
+  return { runner, state: {}, params: {} } as unknown as EvalContext;
+}
+
+function refTerm(refIdx: number, refLen: number, ready = true): RefTerm {
+  const ramp = (base: number, w: number) =>
+    new Float32Array(Array.from({ length: w }, (_v, i) => base + i));
+  return {
+    isReady: () => ready,
+    refIdx,
+    refLen,
+    refRootPos: Array.from({ length: refLen }, (_v, t) => ramp(t * 10, 3)),
+    refRootQuat: Array.from({ length: refLen }, () => new Float32Array([1, 0, 0, 0])),
+    refJointPos: Array.from({ length: refLen }, (_v, t) => ramp(t * 100, 3)),
+  };
+}
+
+describe('DSL arithmetic primitives', () => {
+  it('Div broadcasts a scalar divisor over a vector', () => {
+    expectClose(Primitives.Div([new Float32Array([3, 0, 4]), 2], {}, ctx(null), {}) as Float32Array, [1.5, 0, 2]);
+  });
+  it('Sqrt is elementwise', () => {
+    expectClose(Primitives.Sqrt([new Float32Array([9, 16])], {}, ctx(null), {}) as Float32Array, [3, 4]);
+  });
+  it('Sum reduces a vector to its scalar sum', () => {
+    expect(Primitives.Sum([new Float32Array([3, 0, 4])], {}, ctx(null), {})).toBe(7);
+  });
+  it('Div/Sqrt/Sum compose to L2-normalize (3,0,4 → 0.6,0,0.8)', () => {
+    const v = new Float32Array([3, 0, 4]);
+    const norm = Primitives.Sqrt([Primitives.Sum([Primitives.Mul([v, v], {}, ctx(null), {})], {}, ctx(null), {})], {}, ctx(null), {});
+    expectClose(Primitives.Div([v, norm], {}, ctx(null), {}) as Float32Array, [0.6, 0, 0.8]);
+  });
+});
+
+describe('Slice primitive', () => {
+  it('extracts the contiguous sub-range [start, start+len)', () => {
+    const v = new Float32Array([10, 11, 12, 13, 14, 15]);
+    expectClose(Primitives.Slice([v], { start: 2, len: 3 }, ctx(null), {}) as Float32Array, [12, 13, 14]);
+  });
+  it('zero-pads past the end', () => {
+    expectClose(Primitives.Slice([new Float32Array([1, 2])], { start: 1, len: 3 }, ctx(null), {}) as Float32Array, [2, 0, 0]);
+  });
+});
+
+describe('motion-reference primitives', () => {
+  it('TrackingRefField reads the field at refIdx+step, clamped to the clip', () => {
+    const t = refTerm(2, 5);
+    expectClose(Primitives.TrackingRefField([], { field: 'root_pos', step: 1 }, ctx(t), {}) as Float32Array, t.refRootPos[3]);
+    expectClose(Primitives.TrackingRefField([], { field: 'joint_pos', step: -5 }, ctx(t), {}) as Float32Array, t.refJointPos[0]);
+    expectClose(Primitives.TrackingRefField([], { field: 'root_pos', step: 99 }, ctx(t), {}) as Float32Array, t.refRootPos[4]);
+  });
+  it('TrackingRefField falls back to a finite value when not ready / absent', () => {
+    const nr = ctx(refTerm(0, 4, false));
+    expect(Array.from(Primitives.TrackingRefField([], { field: 'root_quat', step: 0 }, nr, {}) as Float32Array)).toEqual([1, 0, 0, 0]);
+    expect((Primitives.TrackingRefField([], { field: 'root_pos', step: 0 }, nr, {}) as Float32Array).length).toBe(3);
+    // joint width falls back to the policy action count when no clip is loaded.
+    expect((Primitives.TrackingRefField([], { field: 'joint_pos', step: 0 }, ctx(null, 5), {}) as Float32Array).length).toBe(5);
+  });
+  it('TrackingIsReady reflects readiness / presence', () => {
+    expect(Primitives.TrackingIsReady([], {}, ctx(refTerm(0, 4, true)), {})).toBe(1);
+    expect(Primitives.TrackingIsReady([], {}, ctx(refTerm(0, 4, false)), {})).toBe(0);
+    expect(Primitives.TrackingIsReady([], {}, ctx(null), {})).toBe(0);
+  });
+});

package/src/core/dsl/primitives.ts

@@ -71,6 +71,41 @@ function trackingTerm(ctx: EvalContext): TrackingSource | null {
   return isTrackingSource(term) && term.isReady() ? (term as TrackingSource) : null;
 }
 
+/** A motion command exposing reference-trajectory buffers — distinct from the
+ * anchor-based {@link TrackingSource}; built-in and bespoke players declare them. */
+type MotionRefSource = {
+  isReady(): boolean;
+  refRootPos: Float32Array[];
+  refRootQuat: Float32Array[];
+  refJointPos: Float32Array[];
+  refIdx: number;
+  refLen: number;
+};
+
+/** The active motion command carrying reference buffers, or null. */
+function motionRefTerm(ctx: EvalContext): MotionRefSource | null {
+  const term = ctx.runner.getContext()?.commandManager?.getTerm('motion');
+  if (
+    typeof term === 'object' &&
+    term !== null &&
+    'refRootPos' in term &&
+    'refRootQuat' in term &&
+    'refJointPos' in term &&
+    typeof (term as { isReady?: unknown }).isReady === 'function'
+  ) {
+    return term as unknown as MotionRefSource;
+  }
+  return null;
+}
+
+/** Add `step` to `base`, clamped to `[0, len-1]` (matches clampFutureIndices). */
+function clampRefIdx(base: number, step: number, len: number): number {
+  const i = base + step;
+  if (i < 0) return 0;
+  if (i >= len) return Math.max(0, len - 1);
+  return i;
+}
+
 // --- Model name-table decoders + address resolution --------------------------
 
 function decodeNames(
@@ -197,10 +232,20 @@ export const Primitives: Record<string, Primitive> = {
   Add: (inputs) => elementwise(inputs[0], inputs[1], (a, b) => a + b),
   Sub: (inputs) => elementwise(inputs[0], inputs[1], (a, b) => a - b),
   Mul: (inputs) => elementwise(inputs[0], inputs[1], (a, b) => a * b),
+  Div: (inputs) => elementwise(inputs[0], inputs[1], (a, b) => a / b),
   Neg: (inputs) => unary(inputs[0], (a) => -a),
   Abs: (inputs) => unary(inputs[0], Math.abs),
+  Sqrt: (inputs) => unary(inputs[0], Math.sqrt),
   Acos: (inputs) => unary(inputs[0], (a) => Math.acos(Math.max(-1, Math.min(1, a)))),
 
+  // Reduce a vector to the scalar sum of its elements.
+  Sum: (inputs) => {
+    const v = asVec(inputs[0]);
+    let s = 0;
+    for (let i = 0; i < v.length; i++) s += v[i];
+    return s;
+  },
+
   // -- Indexing ------------------------------------------------------------
   Index: (inputs, attrs) => {
     const vec = asVec(inputs[0]);
@@ -208,6 +253,16 @@ export const Primitives: Record<string, Primitive> = {
     return vec[idx] ?? 0;
   },
 
+  // Extract the contiguous sub-range [start, start+len) of a vector.
+  Slice: (inputs, attrs) => {
+    const vec = asVec(inputs[0]);
+    const start = Number(attrs.start ?? 0);
+    const len = Number(attrs.len ?? 0);
+    const out = new Float32Array(len);
+    for (let i = 0; i < len; i++) out[i] = vec[start + i] ?? 0;
+    return out;
+  },
+
   // -- Comparisons (produce Uint8Array for vectors, boolean for scalars) --
   Gt: (inputs) => compare(inputs[0], inputs[1], (a, b) => a > b),
   Lt: (inputs) => compare(inputs[0], inputs[1], (a, b) => a < b),
@@ -437,6 +492,37 @@ export const Primitives: Record<string, Primitive> = {
     return new Float32Array(xquat.slice(id * 4, id * 4 + 4));
   },
 
+  // -- Motion reference trajectory, sampled at a lookahead offset ----------
+  // `field`: 'root_pos' (3) | 'root_quat' (4) | 'joint_pos' (nJoints).
+  // `step`: offset added to refIdx, clamped to [0, refLen-1].
+  // Not ready → finite fallback (identity quat / zeros); pair with
+  // `TrackingIsReady` to zero the term out.
+  TrackingRefField: (_in, attrs, ctx) => {
+    const field = String(attrs.field ?? '');
+    const fallback = () =>
+      field === 'root_quat'
+        ? new Float32Array([1, 0, 0, 0])
+        : field === 'root_pos'
+          ? new Float32Array(3)
+          : new Float32Array(ctx.runner.getNumActions());
+    const term = motionRefTerm(ctx);
+    if (!term || !term.isReady()) return fallback();
+    const buf =
+      field === 'root_pos'
+        ? term.refRootPos
+        : field === 'root_quat'
+          ? term.refRootQuat
+          : term.refJointPos;
+    const frame = buf[clampRefIdx(term.refIdx, Number(attrs.step ?? 0), term.refLen)];
+    return frame ? new Float32Array(frame) : fallback();
+  },
+
+  // 1.0 when a motion reference is loaded and ready, else 0.0.
+  TrackingIsReady: (_in, _attrs, ctx) => {
+    const term = motionRefTerm(ctx);
+    return term && term.isReady() ? 1.0 : 0.0;
+  },
+
   // -- Tracking reference, per body (world frame) --------------------------
   TrackingRefBodyPos: (_in, attrs, ctx) => {
     const term = trackingTerm(ctx);