diff --git a/apps/typegpu-docs/public/assets/genetic-car/car.png b/apps/typegpu-docs/public/assets/genetic-car/car.png new file mode 100644 index 0000000000..569d48f7a8 Binary files /dev/null and b/apps/typegpu-docs/public/assets/genetic-car/car.png differ diff --git a/apps/typegpu-docs/src/examples/algorithms/genetic-racing/ga.ts b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/ga.ts new file mode 100644 index 0000000000..2bf055d1a6 --- /dev/null +++ b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/ga.ts @@ -0,0 +1,320 @@ +import { randf } from '@typegpu/noise'; +import tgpu, { d, std } from 'typegpu'; +import type { TgpuRoot, TgpuUniform } from 'typegpu'; + +export const MAX_POP = 65536; +export const DEFAULT_POP = 8192; + +export const CarState = d.struct({ + position: d.vec2f, + angle: d.f32, + alive: d.u32, + progress: d.f32, + speed: d.f32, + angVel: d.f32, + aliveSteps: d.u32, + stallSteps: d.u32, +}); + +export const FitnessArray = d.arrayOf(d.f32, MAX_POP); + +export const InputLayer = d.struct({ + wA: d.mat4x4f, // inputs[0..3] + wB: d.mat4x4f, // inputs[4..7] + wC: d.mat4x4f, // inputs[8..11] + bias: d.vec4f, +}); + +export const DenseLayer = d.struct({ + w: d.mat4x4f, + bias: d.vec4f, +}); + +export const OutputLayer = d.struct({ + steer: d.vec4f, + throttle: d.vec4f, + bias: d.vec2f, +}); + +export const Genome = d.struct({ + h1: InputLayer, + h2: DenseLayer, + out: OutputLayer, +}); + +export const SimParams = d.struct({ + dt: d.f32, + aspect: d.f32, + generation: d.f32, + population: d.u32, + maxSpeed: d.f32, + accel: d.f32, + turnRate: d.f32, + drag: d.f32, + sensorDistance: d.f32, + mutationRate: d.f32, + mutationStrength: d.f32, + carSize: d.f32, + trackScale: d.f32, + trackLength: d.f32, + spawnX: d.f32, + spawnY: d.f32, + spawnAngle: d.f32, + stepsPerDispatch: d.u32, +}); + +export const CarStateArray = d.arrayOf(CarState, MAX_POP); +export const GenomeArray = d.arrayOf(Genome, MAX_POP); +export const CarStateLayout = d.arrayOf(CarState); + +export const paramsAccess = tgpu.accessor(SimParams); + +const fitLayout = tgpu.bindGroupLayout({ + state: { storage: CarStateArray }, + fitness: { storage: FitnessArray, access: 'mutable' }, +}); + +const initLayout = tgpu.bindGroupLayout({ + state: { storage: CarStateArray, access: 'mutable' }, + genome: { storage: GenomeArray, access: 'mutable' }, +}); + +const evolveLayout = tgpu.bindGroupLayout({ + fitness: { storage: FitnessArray }, + genome: { storage: GenomeArray }, + nextState: { storage: CarStateArray, access: 'mutable' }, + nextGenome: { storage: GenomeArray, access: 'mutable' }, + bestIdx: { storage: d.u32 }, +}); + +const randSignedVec4 = () => { + 'use gpu'; + return (d.vec4f(randf.sample(), randf.sample(), randf.sample(), randf.sample()) * 2 - 1) * 0.8; +}; + +const randSignedMat4x4 = () => { + 'use gpu'; + return d.mat4x4f(randSignedVec4(), randSignedVec4(), randSignedVec4(), randSignedVec4()); +}; + +const makeSpawnState = () => { + 'use gpu'; + const spawn = d.vec2f(paramsAccess.$.spawnX, paramsAccess.$.spawnY) * paramsAccess.$.trackScale; + return CarState({ + position: spawn, + angle: paramsAccess.$.spawnAngle, + speed: 0, + alive: 1, + progress: 0, + angVel: 0, + aliveSteps: 0, + stallSteps: 0, + }); +}; + +const tournamentSelect = () => { + 'use gpu'; + const population = d.f32(paramsAccess.$.population); + let best = d.u32(0); + let bestFitness = d.f32(-1); + for (let j = 0; j < 8; j++) { + const idx = d.u32(randf.sample() * population); + const f = evolveLayout.$.fitness[idx]; + const better = f > bestFitness; + bestFitness = std.select(bestFitness, f, better); + best = std.select(best, idx, better); + } + return best; +}; + +const evolveVec = (a: T, b: T): T => { + 'use gpu'; + const strength = paramsAccess.$.mutationStrength; + const crossed = std.select(a, b, randf.sample() > 0.5); + const doMutate = randf.sample() < paramsAccess.$.mutationRate; + if (a.kind === 'vec2f') { + const delta = d.vec2f(randf.normal(0, strength), randf.normal(0, strength)); + return ((crossed as d.v2f) + std.select(d.vec2f(0), delta, doMutate)) as T; + } else { + const delta = d.vec4f( + randf.normal(0, strength), + randf.normal(0, strength), + randf.normal(0, strength), + randf.normal(0, strength), + ); + return ((crossed as d.v4f) + std.select(d.vec4f(0), delta, doMutate)) as T; + } +}; + +const evolveMat4x4 = (a: d.m4x4f, b: d.m4x4f) => { + 'use gpu'; + return d.mat4x4f( + evolveVec(a.columns[0], b.columns[0]), + evolveVec(a.columns[1], b.columns[1]), + evolveVec(a.columns[2], b.columns[2]), + evolveVec(a.columns[3], b.columns[3]), + ); +}; + +const evolveInputLayer = (a: d.InferGPU, b: d.InferGPU) => { + 'use gpu'; + return InputLayer({ + wA: evolveMat4x4(a.wA, b.wA), + wB: evolveMat4x4(a.wB, b.wB), + wC: evolveMat4x4(a.wC, b.wC), + bias: evolveVec(a.bias, b.bias), + }); +}; + +const evolveDenseLayer = (a: d.InferGPU, b: d.InferGPU) => { + 'use gpu'; + return DenseLayer({ w: evolveMat4x4(a.w, b.w), bias: evolveVec(a.bias, b.bias) }); +}; + +const evolveOutputLayer = ( + a: d.InferGPU, + b: d.InferGPU, +) => { + 'use gpu'; + return OutputLayer({ + steer: evolveVec(a.steer, b.steer), + throttle: evolveVec(a.throttle, b.throttle), + bias: evolveVec(a.bias, b.bias), + }); +}; + +const fitShader = (i: number) => { + 'use gpu'; + if (d.u32(i) >= paramsAccess.$.population) { + return; + } + const s = CarState(fitLayout.$.state[i]); + fitLayout.$.fitness[i] = s.progress * 10 + d.f32(s.aliveSteps) * 0.003; +}; + +const initShader = (i: number) => { + 'use gpu'; + if (d.u32(i) >= paramsAccess.$.population) { + return; + } + randf.seed2(d.vec2f(d.f32(i) + 1, paramsAccess.$.generation + 11)); + + initLayout.$.genome[i] = Genome({ + h1: { + wA: randSignedMat4x4(), + wB: randSignedMat4x4(), + wC: randSignedMat4x4(), + bias: d.vec4f(), + }, + h2: { w: randSignedMat4x4(), bias: d.vec4f() }, + out: { steer: randSignedVec4(), throttle: randSignedVec4(), bias: d.vec2f() }, + }); + initLayout.$.state[i] = makeSpawnState(); +}; + +const evolveShader = (i: number) => { + 'use gpu'; + if (d.u32(i) >= paramsAccess.$.population) { + return; + } + + // Elitism: champion always lives at index 0, copied unchanged + if (d.u32(i) === 0) { + evolveLayout.$.nextGenome[0] = Genome(evolveLayout.$.genome[evolveLayout.$.bestIdx]); + evolveLayout.$.nextState[0] = makeSpawnState(); + return; + } + + randf.seed2(d.vec2f(d.f32(i) + 3, paramsAccess.$.generation + 19)); + + const parentA = Genome(evolveLayout.$.genome[tournamentSelect()]); + const parentB = Genome(evolveLayout.$.genome[tournamentSelect()]); + + evolveLayout.$.nextGenome[i] = Genome({ + h1: evolveInputLayer(parentA.h1, parentB.h1), + h2: evolveDenseLayer(parentA.h2, parentB.h2), + out: evolveOutputLayer(parentA.out, parentB.out), + }); + + evolveLayout.$.nextState[i] = makeSpawnState(); +}; + +export function createGeneticPopulation(root: TgpuRoot, params: TgpuUniform) { + const stateBuffers = [0, 1].map(() => + root.createBuffer(CarStateArray).$usage('storage', 'vertex'), + ); + const genomeBuffers = [0, 1].map(() => root.createBuffer(GenomeArray).$usage('storage')); + const fitnessBuffer = root.createBuffer(FitnessArray).$usage('storage'); + const bestIdxBuffer = root.createBuffer(d.u32).$usage('storage'); + + const initBindGroups = [0, 1].map((i) => + root.createBindGroup(initLayout, { + state: stateBuffers[i], + genome: genomeBuffers[i], + }), + ); + + const fitBindGroups = [0, 1].map((i) => + root.createBindGroup(fitLayout, { + state: stateBuffers[i], + fitness: fitnessBuffer, + }), + ); + + const evolveBindGroups = [0, 1].map((i) => + root.createBindGroup(evolveLayout, { + fitness: fitnessBuffer, + genome: genomeBuffers[i], + nextState: stateBuffers[1 - i], + nextGenome: genomeBuffers[1 - i], + bestIdx: bestIdxBuffer, + }), + ); + + const initPipeline = root.with(paramsAccess, params).createGuardedComputePipeline(initShader); + const fitPipeline = root.with(paramsAccess, params).createGuardedComputePipeline(fitShader); + const evolvePipeline = root.with(paramsAccess, params).createGuardedComputePipeline(evolveShader); + + let current = 0; + let generation = 0; + + return { + stateBuffers, + genomeBuffers, + fitnessBuffer, + bestIdxBuffer, + get current() { + return current; + }, + get generation() { + return generation; + }, + get currentStateBuffer() { + return stateBuffers[current]; + }, + get currentGenomeBuffer() { + return genomeBuffers[current]; + }, + + init() { + current = 0; + generation = 0; + initPipeline.with(initBindGroups[0]).dispatchThreads(MAX_POP); + initPipeline.with(initBindGroups[1]).dispatchThreads(MAX_POP); + }, + + reinitCurrent(population: number) { + initPipeline.with(initBindGroups[current]).dispatchThreads(population); + }, + + precomputeFitness(population: number) { + fitPipeline.with(fitBindGroups[current]).dispatchThreads(population); + }, + + evolve(population: number) { + evolvePipeline.with(evolveBindGroups[current]).dispatchThreads(population); + current = 1 - current; + generation++; + }, + }; +} diff --git a/apps/typegpu-docs/src/examples/algorithms/genetic-racing/index.html b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/index.html new file mode 100644 index 0000000000..e049138b14 --- /dev/null +++ b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/index.html @@ -0,0 +1,14 @@ + +
+ diff --git a/apps/typegpu-docs/src/examples/algorithms/genetic-racing/index.ts b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/index.ts new file mode 100644 index 0000000000..93558cf937 --- /dev/null +++ b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/index.ts @@ -0,0 +1,611 @@ +import tgpu, { common, d, std } from 'typegpu'; +import { defineControls } from '../../common/defineControls.ts'; +import { generateGridTrack, type TrackResult } from './track.ts'; +import { + CarState, + CarStateArray, + CarStateLayout, + DEFAULT_POP, + FitnessArray, + Genome, + GenomeArray, + MAX_POP, + SimParams, + createGeneticPopulation, +} from './ga.ts'; + +const DEG_90 = Math.PI / 2; +const DEG_60 = Math.PI / 3; +const DEG_30 = Math.PI / 6; + +const root = await tgpu.init(); +const canvas = document.querySelector('canvas') as HTMLCanvasElement; +const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); + +const STEPS_PER_DISPATCH = 32; + +const BASE_SPATIAL_PARAMS = { + maxSpeed: 1.6, + accel: 0.2, + turnRate: 5.5, + drag: 0.3, + sensorDistance: 0.28, + carSize: 0.02, +}; + +const params = root.createUniform(SimParams, { + dt: 1 / 120, + aspect: 1, + generation: 0, + population: DEFAULT_POP, + mutationRate: 0.05, + mutationStrength: 0.15, + trackScale: 0.9, + trackLength: 1, + spawnX: 0, + spawnY: 0, + spawnAngle: 0, + stepsPerDispatch: STEPS_PER_DISPATCH, + ...BASE_SPATIAL_PARAMS, +}); + +const ga = createGeneticPopulation(root, params); + +const trackTexture = root['~unstable'] + .createTexture({ size: [512, 512], format: 'rgba8unorm' }) + .$usage('render', 'sampled'); +const trackView = trackTexture.createView(); + +const carBitmap = await fetch('/TypeGPU/assets/genetic-car/car.png') + .then((r) => r.blob()) + .then(createImageBitmap); + +const carSpriteTexture = root['~unstable'] + .createTexture({ + size: [carBitmap.width / 2, carBitmap.height / 2], + format: 'rgba8unorm', + }) + .$usage('sampled', 'render'); +carSpriteTexture.write(carBitmap); +const carSpriteView = carSpriteTexture.createView(); + +const linearSampler = root['~unstable'].createSampler({ + magFilter: 'linear', + minFilter: 'linear', +}); +const nearestSampler = root['~unstable'].createSampler({ + magFilter: 'nearest', + minFilter: 'nearest', +}); + +const toTrackSpace = (p: d.v2f) => { + 'use gpu'; + return p / params.$.trackScale; +}; + +const toTrackUV = (p: d.v2f) => { + 'use gpu'; + const uvBase = (toTrackSpace(p) + 1) * 0.5; + return d.vec2f(uvBase.x, 1 - uvBase.y); +}; + +const sampleTrack = (p: d.v2f, sampler: d.sampler) => { + 'use gpu'; + const sample = std.textureSampleLevel(trackView.$, sampler, toTrackUV(p), 0); + return d.vec3f(sample.xy * 2 - 1, sample.z); +}; + +const rotate = (v: d.v2f, angle: number) => { + 'use gpu'; + const c = std.cos(angle); + const s = std.sin(angle); + return d.vec2f(v.x * c - v.y * s, v.x * s + v.y * c); +}; + +const isOnTrack = (pos: d.v2f) => { + 'use gpu'; + return sampleTrack(pos, nearestSampler.$).z > 0.5; +}; + +const trackCross = (forward: d.v2f, pos: d.v2f) => { + 'use gpu'; + const t = sampleTrack(pos, nearestSampler.$); + return forward.x * t.y - forward.y * t.x; +}; + +const senseRaycast = (pos: d.v2f, angle: number, offset: number) => { + 'use gpu'; + const dir = d.vec2f(std.cos(angle + offset), std.sin(angle + offset)); + let hitT = d.f32(1); + for (const step of tgpu.unroll([1, 2, 3, 4, 5, 6, 7, 8])) { + const t = d.f32(step / 8); + const samplePos = pos + dir * t * params.$.sensorDistance; + const s = sampleTrack(samplePos, nearestSampler.$); + hitT = std.select(hitT, std.select(t, hitT, hitT < t), s.z < 0.5); + } + return hitT; +}; + +const evalNetwork = (genome: d.InferGPU, a: d.v4f, b: d.v4f, c: d.v4f) => { + 'use gpu'; + const h1 = std.tanh( + std.transpose(genome.h1.wA) * a + + std.transpose(genome.h1.wB) * b + + std.transpose(genome.h1.wC) * c + + genome.h1.bias, + ); + const h2 = std.tanh(std.transpose(genome.h2.w) * h1 + genome.h2.bias); + return std.clamp( + d.vec2f(std.dot(genome.out.steer, h2), std.dot(genome.out.throttle, h2)) + genome.out.bias, + d.vec2f(-1), + d.vec2f(1), + ); +}; + +const simLayout = tgpu.bindGroupLayout({ + state: { storage: CarStateArray, access: 'mutable' }, + genome: { storage: GenomeArray }, +}); + +const simBindGroups = [0, 1].map((i) => + root.createBindGroup(simLayout, { + state: ga.stateBuffers[i], + genome: ga.genomeBuffers[i], + }), +); + +const simulatePipeline = root.createGuardedComputePipeline((i) => { + 'use gpu'; + if (d.u32(i) >= params.$.population) { + return; + } + + const genome = Genome(simLayout.$.genome[i]); + const initCar = CarState(simLayout.$.state[i]); + + let curPosition = d.vec2f(initCar.position); + let curAngle = initCar.angle; + let curSpeed = initCar.speed; + let curAlive = initCar.alive; + let curProgress = initCar.progress; + let curAngVel = initCar.angVel; + let curAliveSteps = initCar.aliveSteps; + let curStallSteps = initCar.stallSteps; + + for (let s = d.u32(0); s < params.$.stepsPerDispatch; s++) { + if (curAlive === 0) { + break; + } + + const carForward = d.vec2f(std.cos(curAngle), std.sin(curAngle)); + const aheadPos = curPosition + carForward * params.$.sensorDistance; + + const inputs4 = d.vec4f( + senseRaycast(curPosition, curAngle, DEG_60), + senseRaycast(curPosition, curAngle, DEG_30), + senseRaycast(curPosition, curAngle, 0), + senseRaycast(curPosition, curAngle, -DEG_30), + ); + const inputsB = d.vec4f( + senseRaycast(curPosition, curAngle, -DEG_60), + curSpeed / params.$.maxSpeed, + std.dot(carForward, sampleTrack(curPosition, nearestSampler.$).xy), + trackCross(carForward, aheadPos), + ); + const inputsC = d.vec4f( + curAngVel / params.$.turnRate, + senseRaycast(curPosition, curAngle, DEG_90), + senseRaycast(curPosition, curAngle, -DEG_90), + trackCross(carForward, curPosition + carForward * params.$.sensorDistance * 2), + ); + + const control = evalNetwork(genome, inputs4, inputsB, inputsC); + const steer = control.x; + const throttle = control.y; + + let speed = curSpeed + throttle * params.$.accel * params.$.dt; + speed = speed * (1 - params.$.drag * speed * params.$.dt); + speed = std.clamp(speed, 0, params.$.maxSpeed); + + const slowThreshold = params.$.maxSpeed * 0.04; + const canTurn = speed > slowThreshold; + const normSpeed = speed / params.$.maxSpeed; + const turnFactor = (1 - normSpeed) * (1 - normSpeed); + const targetAngVel = std.select(0, steer * params.$.turnRate * turnFactor, canTurn); + const angVel = curAngVel * 0.75 + targetAngVel * 0.25; + const angle = curAngle + angVel * params.$.dt; + + const dir = d.vec2f(std.cos(angle), std.sin(angle)); + const position = curPosition + dir * speed * params.$.dt; + const stepVec = position - curPosition; + + const stallSteps = std.select(d.u32(0), curStallSteps + 1, speed < slowThreshold); + const trackEnd = sampleTrack(position, nearestSampler.$); + const onTrack = + stallSteps < 120 && + trackEnd.z > 0.5 && + isOnTrack(curPosition + stepVec * 0.33) && + isOnTrack(curPosition + stepVec * 0.66); + + const alive = std.select(d.u32(0), d.u32(1), onTrack); + const forward = std.dot(dir, trackEnd.xy); + const lapLength = params.$.trackLength * params.$.trackScale; + + curPosition = std.select(curPosition, position, onTrack); + curAngle = std.select(curAngle, angle, onTrack); + curSpeed = std.select(0, speed, onTrack); + curAlive = alive; + curProgress = + curProgress + (speed * std.max(0, forward) * params.$.dt * d.f32(alive)) / lapLength; + curAngVel = std.select(0, angVel, onTrack); + curAliveSteps = curAliveSteps + 1; + curStallSteps = stallSteps; + } + + simLayout.$.state[i] = CarState({ + position: curPosition, + angle: curAngle, + speed: curSpeed, + alive: curAlive, + progress: curProgress, + angVel: curAngVel, + aliveSteps: curAliveSteps, + stallSteps: curStallSteps, + }); +}); + +// upper 16 bits = quantized fitness [0,65535], lower 16 bits = car index +const reductionPackedBuffer = root.createBuffer(d.atomic(d.u32), 0).$usage('storage'); +const bestFitnessBuffer = root.createBuffer(d.f32).$usage('storage'); + +const reductionLayout = tgpu.bindGroupLayout({ + fitness: { storage: FitnessArray }, + genome: { storage: GenomeArray }, + packed: { storage: d.atomic(d.u32), access: 'mutable' }, + bestIdx: { storage: d.u32, access: 'mutable' }, + bestFitness: { storage: d.f32, access: 'mutable' }, +}); + +const reductionBindGroups = [0, 1].map((i) => + root.createBindGroup(reductionLayout, { + fitness: ga.fitnessBuffer, + genome: ga.genomeBuffers[i], + bestIdx: ga.bestIdxBuffer, + packed: reductionPackedBuffer, + bestFitness: bestFitnessBuffer, + }), +); + +const reductionPipeline = root.createGuardedComputePipeline((i) => { + 'use gpu'; + if (d.u32(i) >= params.$.population) { + return; + } + const fitness = reductionLayout.$.fitness[i]; + const quantized = d.u32(std.clamp(fitness / 64, 0, 1) * 65535); + const packed = (quantized << 16) | (d.u32(i) & 0xffff); + std.atomicMax(reductionLayout.$.packed, packed); +}); + +const finalizeReductionPipeline = root.createGuardedComputePipeline(() => { + 'use gpu'; + const packed = std.atomicLoad(reductionLayout.$.packed); + reductionLayout.$.bestIdx = packed & 0xffff; + reductionLayout.$.bestFitness = (d.f32(packed >> 16) / 65535) * 64; +}); + +const colors = { + grass: tgpu.const(d.vec3f, d.vec3f(0.05, 0.06, 0.08)), + road: tgpu.const(d.vec3f, d.vec3f(0.14, 0.16, 0.2)), + paint: tgpu.const(d.vec3f, d.vec3f(0.2, 0.22, 0.3)), +}; + +const trackFragment = tgpu.fragmentFn({ + in: { uv: d.vec2f }, + out: d.vec4f, +})(({ uv }) => { + 'use gpu'; + const p = d.vec2f((uv.x * 2 - 1) * params.$.aspect, 1 - uv.y * 2); + const sample = sampleTrack(p, linearSampler.$); + + const mask = sample.z; + const color = std.mix(colors.grass.$, colors.road.$, mask); + const edge = 1 - std.smoothstep(0.6, 0.95, mask); + const painted = color + colors.paint.$ * edge * mask; + + return d.vec4f(painted, 1); +}); + +const trackPipeline = root.createRenderPipeline({ + vertex: common.fullScreenTriangle, + fragment: trackFragment, +}); + +const carQuad = tgpu.const(d.arrayOf(d.vec4f, 4), [ + d.vec4f(-1, -1, 0, 1), + d.vec4f(1, -1, 0, 0), + d.vec4f(-1, 1, 1, 1), + d.vec4f(1, 1, 1, 0), +]); + +const carVertex = tgpu.vertexFn({ + in: { + vertexIndex: d.builtin.vertexIndex, + position: d.vec2f, + angle: d.f32, + alive: d.u32, + progress: d.f32, + }, + out: { pos: d.builtin.position, uv: d.vec2f, isAlive: d.f32, progress: d.f32 }, +})((input) => { + 'use gpu'; + const q = carQuad.$[input.vertexIndex]; + const localPos = d.vec2f(q.x, q.y * 0.5) * params.$.carSize; + const rotated = rotate(localPos, input.angle); + const worldPos = rotated + input.position; + const pos = d.vec4f(worldPos.x / params.$.aspect, worldPos.y, 0, 1); + const isAlive = std.select(0, d.f32(1), input.alive === 1); + return { pos, uv: q.zw, isAlive, progress: input.progress }; +}); + +const carFragment = tgpu.fragmentFn({ + in: { uv: d.vec2f, isAlive: d.f32, progress: d.f32 }, + out: d.vec4f, +})(({ uv, isAlive, progress }) => { + 'use gpu'; + const sample = std.textureSampleLevel(carSpriteView.$, linearSampler.$, uv, 0); + const t = std.smoothstep(0, 1, progress); + const baseTint = std.mix(d.vec3f(0.4, 0.6, 1.0), d.vec3f(1.0, 0.85, 0.15), t); + const lapAccent = std.smoothstep(1, 10, progress); + const tint = std.mix(baseTint, d.vec3f(0.15, 1.0, 0.35), lapAccent); + const lum = std.dot(sample.xyz, d.vec3f(0.299, 0.587, 0.114)); + const rgb = std.mix(d.vec3f(lum) * 0.4, sample.xyz * tint, isAlive); + const a = sample.w * std.mix(0.45, 1, isAlive); + return d.vec4f(rgb * a, a); +}); + +const instanceLayout = tgpu.vertexLayout(CarStateLayout, 'instance'); + +const carPipeline = root.createRenderPipeline({ + attribs: instanceLayout.attrib, + vertex: carVertex, + fragment: carFragment, + primitive: { topology: 'triangle-strip' }, + targets: { + blend: { + color: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha', operation: 'add' }, + alpha: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha', operation: 'add' }, + }, + }, +}); + +let steps = 0; +let stepsPerFrame = STEPS_PER_DISPATCH; +let stepsPerGeneration = 2048; +let paused = false; +let lastAspect = 1; +let population = DEFAULT_POP; +let rafHandle = 0; +let pendingEvolve = false; +let showBestOnly = false; +let displayedBestFitness = 0; + +const statsDiv = document.querySelector('.stats') as HTMLDivElement; + +function updateAspect() { + if (!canvas.width || !canvas.height) { + return; + } + const nextAspect = canvas.width / canvas.height; + if (Math.abs(nextAspect - lastAspect) < 0.001) { + return; + } + lastAspect = nextAspect; + params.writePartial({ aspect: nextAspect }); +} + +function updatePopulation(nextPopulation: number) { + const clamped = Math.max(128, Math.min(MAX_POP, Math.floor(nextPopulation))); + if (clamped === population) { + return; + } + population = clamped; + params.writePartial({ population: clamped }); + ga.reinitCurrent(population); +} + +function frame() { + updateAspect(); + + if (!paused) { + if (pendingEvolve) { + ga.evolve(population); + steps = 0; + params.writePartial({ generation: ga.generation }); + pendingEvolve = false; + } + + const stepsToRun = Math.min(stepsPerFrame, stepsPerGeneration - steps); + if (stepsToRun <= 0) { + pendingEvolve = true; + } else { + const innerSteps = Math.min(stepsToRun, STEPS_PER_DISPATCH); + params.writePartial({ stepsPerDispatch: innerSteps }); + const dispatchCount = Math.ceil(stepsToRun / innerSteps); + + const simEncoder = root.device.createCommandEncoder(); + const encoderPipeline = simulatePipeline.with(simBindGroups[ga.current]).with(simEncoder); + for (let dispatch = 0; dispatch < dispatchCount; dispatch++) { + encoderPipeline.dispatchThreads(population); + } + root.device.queue.submit([simEncoder.finish()]); + + steps += dispatchCount * innerSteps; + } + + if (steps >= stepsPerGeneration) { + pendingEvolve = true; + ga.precomputeFitness(population); + const bg = reductionBindGroups[ga.current]; + + const reductionEncoder = root.device.createCommandEncoder(); + reductionEncoder.clearBuffer(root.unwrap(reductionPackedBuffer)); + reductionPipeline.with(bg).with(reductionEncoder).dispatchThreads(population); + finalizeReductionPipeline.with(bg).with(reductionEncoder).dispatchThreads(); + root.device.queue.submit([reductionEncoder.finish()]); + + void bestFitnessBuffer.read().then((fitness) => { + displayedBestFitness = fitness; + }); + } + } + + const genStr = String(ga.generation).padStart(5); + const stepStr = String(steps).padStart(String(stepsPerGeneration).length); + const bestStr = displayedBestFitness.toFixed(2).padStart(6); + const saturatedNote = displayedBestFitness >= 64 ? ' (saturated)' : ''; + statsDiv.textContent = `Gen ${genStr} Step ${stepStr}/${stepsPerGeneration} Pop ${population} Best ${bestStr}${saturatedNote}`; + + trackPipeline.withColorAttachment({ view: context, clearValue: [0.04, 0.05, 0.07, 1] }).draw(3); + + carPipeline + .withColorAttachment({ view: context, loadOp: 'load', storeOp: 'store' }) + .with(instanceLayout, ga.currentStateBuffer) + .draw(4, showBestOnly ? 1 : population); + + rafHandle = requestAnimationFrame(frame); +} + +function applyTrack(result: TrackResult) { + trackTexture.write( + new ImageData(new Uint8ClampedArray(result.data), result.width, result.height), + ); + params.writePartial({ + spawnX: result.spawn.position[0], + spawnY: result.spawn.position[1], + spawnAngle: result.spawn.angle, + trackLength: result.trackLength, + }); +} + +function applyGridSize(W: number, H: number) { + const scale = 5 / Math.max(W, H); + params.writePartial( + Object.fromEntries( + Object.entries(BASE_SPATIAL_PARAMS).map(([k, v]) => [k, v * scale]), + ) as typeof BASE_SPATIAL_PARAMS, + ); +} + +const GRID_SIZES: Record = { + S: [5, 4], + M: [8, 6], + L: [10, 9], + XL: [14, 12], +}; + +let trackSeed = (Math.random() * 100_000) | 0; +let gridSizeKey = 'S'; + +function startSimulation() { + steps = 0; + pendingEvolve = false; + displayedBestFitness = 0; + params.writePartial({ generation: 0 }); + ga.init(); + + updateAspect(); + updatePopulation(population); + cancelAnimationFrame(rafHandle); + rafHandle = requestAnimationFrame(frame); +} + +function newTrack() { + trackSeed = (Math.random() * 100_000) | 0; + const [W, H] = GRID_SIZES[gridSizeKey]; + applyGridSize(W, H); + applyTrack(generateGridTrack(trackSeed, W, H)); + startSimulation(); +} + +applyGridSize(...GRID_SIZES[gridSizeKey]); +applyTrack(generateGridTrack(trackSeed, ...GRID_SIZES[gridSizeKey])); +startSimulation(); + +// #region Example controls & Cleanup + +export const controls = defineControls({ + 'New Track': { onButtonClick: newTrack }, + 'Grid size': { + initial: 'S', + options: ['S', 'M', 'L', 'XL'], + onSelectChange: (value: string) => { + gridSizeKey = value; + newTrack(); + }, + }, + Pause: { + initial: false, + onToggleChange: (value: boolean) => { + paused = value; + }, + }, + 'Best car only': { + initial: false, + onToggleChange: (value: boolean) => { + showBestOnly = value; + }, + }, + 'Steps per frame': { + initial: stepsPerFrame, + min: 1, + max: 8192, + step: 1, + onSliderChange: (value: number) => { + stepsPerFrame = value; + }, + }, + 'Steps per generation': { + initial: stepsPerGeneration, + min: 120, + max: 9600, + step: 60, + onSliderChange: (value: number) => { + stepsPerGeneration = value; + }, + }, + Population: { + initial: population, + min: 256, + max: MAX_POP, + step: 256, + onSliderChange: (value: number) => { + updatePopulation(value); + }, + }, + 'Mutation rate': { + initial: 0.05, + min: 0, + max: 0.4, + step: 0.005, + onSliderChange: (value: number) => { + params.writePartial({ mutationRate: value }); + }, + }, + 'Mutation strength': { + initial: 0.15, + min: 0.01, + max: 0.8, + step: 0.01, + onSliderChange: (value: number) => { + params.writePartial({ mutationStrength: value }); + }, + }, +}); + +export function onCleanup() { + cancelAnimationFrame(rafHandle); + root.destroy(); +} + +// #endregion diff --git a/apps/typegpu-docs/src/examples/algorithms/genetic-racing/meta.json b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/meta.json new file mode 100644 index 0000000000..11aab6386e --- /dev/null +++ b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/meta.json @@ -0,0 +1,6 @@ +{ + "title": "Genetic Racing", + "category": "algorithms", + "tags": ["genetic algorithm", "ai"], + "coolFactor": 7 +} diff --git a/apps/typegpu-docs/src/examples/algorithms/genetic-racing/thumbnail.png b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/thumbnail.png new file mode 100644 index 0000000000..59d6d13193 Binary files /dev/null and b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/thumbnail.png differ diff --git a/apps/typegpu-docs/src/examples/algorithms/genetic-racing/track.ts b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/track.ts new file mode 100644 index 0000000000..f569a0d2e2 --- /dev/null +++ b/apps/typegpu-docs/src/examples/algorithms/genetic-racing/track.ts @@ -0,0 +1,244 @@ +import { std } from 'typegpu'; + +export type TrackResult = { + width: number; + height: number; + data: Uint8ClampedArray; + spawn: { position: [number, number]; angle: number }; + trackLength: number; +}; + +type Pt = { x: number; y: number }; + +function mulberry32(seed: number) { + let t = seed >>> 0; + return () => { + t += 0x6d2b79f5; + let r = t; + r = Math.imul(r ^ (r >>> 15), r | 1); + r ^= r + Math.imul(r ^ (r >>> 7), r | 61); + return ((r ^ (r >>> 14)) >>> 0) / 4294967296; + }; +} + +function findBestSpawnIndex(pts: Pt[], windowSize: number): number { + const n = pts.length; + const curv = new Float64Array(n); + for (let i = 0; i < n; i++) { + const prev = pts[(i - 1 + n) % n]; + const curr = pts[i]; + const next = pts[(i + 1) % n]; + + const t1x = curr.x - prev.x; + const t1y = curr.y - prev.y; + const t1len = Math.hypot(t1x, t1y) || 1; + + const t2x = next.x - curr.x; + const t2y = next.y - curr.y; + const t2len = Math.hypot(t2x, t2y) || 1; + + curv[i] = Math.abs((t1x / t1len) * (t2y / t2len) - (t1y / t1len) * (t2x / t2len)); + } + + let windowSum = 0; + for (let i = 0; i < windowSize; i++) windowSum += curv[i]; + + let minSum = windowSum; + let bestStart = 0; + + for (let i = 1; i < n; i++) { + windowSum += curv[(i + windowSize - 1) % n] - curv[i - 1]; + if (windowSum < minSum) { + minSum = windowSum; + bestStart = i; + } + } + + return (bestStart + Math.floor(windowSize / 2)) % n; +} + +function catmullRomResample(points: Pt[], numSamples: number): Pt[] { + const n = points.length; + const result: Pt[] = []; + + for (let s = 0; s < numSamples; s++) { + const tTotal = s / numSamples; + const seg = Math.floor(tTotal * n); + const t = tTotal * n - seg; + + const i0 = (seg - 1 + n) % n; + const i1 = seg; + const i2 = (seg + 1) % n; + const i3 = (seg + 2) % n; + + const p0 = points[i0]; + const p1 = points[i1]; + const p2 = points[i2]; + const p3 = points[i3]; + + const t2 = t * t; + const t3 = t2 * t; + + result.push({ + x: + 0.5 * + (2 * p1.x + + (-p0.x + p2.x) * t + + (2 * p0.x - 5 * p1.x + 4 * p2.x - p3.x) * t2 + + (-p0.x + 3 * p1.x - 3 * p2.x + p3.x) * t3), + y: + 0.5 * + (2 * p1.y + + (-p0.y + p2.y) * t + + (2 * p0.y - 5 * p1.y + 4 * p2.y - p3.y) * t2 + + (-p0.y + 3 * p1.y - 3 * p2.y + p3.y) * t3), + }); + } + + return result; +} + +function buildTrackTexture( + splinedPath: Pt[], + textureSize: number, + trackHalfWidth: number, + spawnIdx?: number, +): TrackResult { + const feather = trackHalfWidth * 0.4; + + const segments = splinedPath.map((cell, idx) => { + const next = splinedPath[(idx + 1) % splinedPath.length]; + const dx = next.x - cell.x; + const dy = next.y - cell.y; + const len = Math.hypot(dx, dy) || 1; + return { + ax: cell.x, + ay: cell.y, + ux: dx, + uy: dy, + dx: dx / len, + dy: dy / len, + len, + len2: dx * dx + dy * dy, + }; + }); + + let totalLen = 0; + for (const segment of segments) { + totalLen += segment.len; + } + + const data = new Uint8ClampedArray(textureSize * textureSize * 4); + + for (let y = 0; y < textureSize; y++) { + const ty = 1 - (y / (textureSize - 1)) * 2; + for (let x = 0; x < textureSize; x++) { + const tx = (x / (textureSize - 1)) * 2 - 1; + + let minSegDist = Infinity; + let dir = { x: 1, y: 0 }; + for (const segment of segments) { + const px = tx - segment.ax; + const py = ty - segment.ay; + const t = Math.max(0, Math.min(1, (px * segment.ux + py * segment.uy) / segment.len2)); + const cx = segment.ax + segment.ux * t; + const cy = segment.ay + segment.uy * t; + const dx = tx - cx; + const dy = ty - cy; + const dist = dx * dx + dy * dy; + if (dist < minSegDist) { + minSegDist = dist; + dir = { x: segment.dx, y: segment.dy }; + } + } + + const dist = Math.sqrt(minSegDist); + const mask = 1 - std.smoothstep(trackHalfWidth, trackHalfWidth + feather, dist); + + const idx = (y * textureSize + x) * 4; + data[idx] = Math.round((dir.x * 0.5 + 0.5) * 255); + data[idx + 1] = Math.round((dir.y * 0.5 + 0.5) * 255); + data[idx + 2] = Math.round(mask * 255); + data[idx + 3] = 255; + } + } + + const resolvedSpawnIdx = spawnIdx ?? findBestSpawnIndex(splinedPath, 25); + const spawnStart = splinedPath[resolvedSpawnIdx]; + const spawnNext = splinedPath[(resolvedSpawnIdx + 1) % splinedPath.length]; + const spawnAngle = Math.atan2(spawnNext.y - spawnStart.y, spawnNext.x - spawnStart.x); + + return { + width: textureSize, + height: textureSize, + data, + spawn: { + position: [spawnStart.x, spawnStart.y], + angle: spawnAngle, + }, + trackLength: totalLen, + }; +} + +function generateLoopPath(W: number, H: number, rand: () => number): number[] { + const minLength = Math.max(8, Math.floor(W * H * 0.4)); + const path: number[] = [0]; + const inPath = new Map([[0, 0]]); + + for (let attempts = 0; attempts < 2_000_000; attempts++) { + const current = path[path.length - 1]; + const x = current % W, + y = Math.floor(current / W); + + const neighbors: number[] = []; + if (x > 0) neighbors.push(current - 1); + if (x < W - 1) neighbors.push(current + 1); + if (y > 0) neighbors.push(current - W); + if (y < H - 1) neighbors.push(current + W); + + const next = neighbors[Math.floor(rand() * neighbors.length)]; + + if (next === 0 && path.length >= minLength) { + return path; // closed loop found + } + + const existingIdx = inPath.get(next); + if (existingIdx !== undefined) { + // Erase the loop back to where 'next' was first visited + for (let i = existingIdx + 1; i < path.length; i++) { + inPath.delete(path[i]); + } + path.length = existingIdx + 1; + } else { + inPath.set(next, path.length); + path.push(next); + } + } + + const perimeter: number[] = []; + for (let x = 0; x < W; x++) perimeter.push(x); + for (let y = 1; y < H; y++) perimeter.push(y * W + W - 1); + for (let x = W - 2; x >= 0; x--) perimeter.push((H - 1) * W + x); + for (let y = H - 2; y > 0; y--) perimeter.push(y * W); + return perimeter; +} + +export function generateGridTrack(seed: number, W = 5, H = 4, textureSize = 512): TrackResult { + const rand = mulberry32(seed); + const cellPath = generateLoopPath(W, H, rand); + + const controlPoints = cellPath.map((cell) => { + const col = cell % W, + row = Math.floor(cell / W); + const cx = -0.8 + ((col + 0.5) * 1.6) / W; + const cy = 0.8 - ((row + 0.5) * 1.6) / H; + const jx = (rand() * 2 - 1) * 0.06 * (1.6 / W); + const jy = (rand() * 2 - 1) * 0.06 * (1.6 / H); + return { x: cx + jx, y: cy + jy }; + }); + + const trackHalfWidth = 0.172 * (1.6 / Math.max(W, H)); + const numSamples = Math.max(120, cellPath.length * 6); + const splinedPath = catmullRomResample(controlPoints, numSamples); + return buildTrackTexture(splinedPath, textureSize, trackHalfWidth); +}