SirStone
0d35646dc9
- actorForward: deterministic param, uses mean-only when PPOB_EVAL_ONLY=1 (eval was adding unit Gaussian noise to every action — unreliable scores) - warm_start.py: log_std initialized to -1.0 (std≈0.37) instead of copying snapshot values (were 2.27-4.68 → std 9-108, completely drowning signal) - training.env: LOG_STD_CEILING 0.0→-0.5 (cap exploration at std≈0.6)
87 lines
3.5 KiB
Nim
87 lines
3.5 KiB
Nim
## network.nim — MLP and ActorCritic forward pass (inference only, no autograd).
|
|
|
|
import arraymancer
|
|
import std/[math, random]
|
|
|
|
const
|
|
STATE_DIM* = 57
|
|
ACTION_DIM* = 6
|
|
|
|
var
|
|
logStdFloor*: float32 = -3.0'f32 # overridden by PPOB_LOG_STD_FLOOR
|
|
logStdCeiling*: float32 = 0.5'f32 # overridden by PPOB_LOG_STD_CEILING
|
|
initialLogStd*: float32 = 0.0'f32 # overridden by PPOB_INITIAL_LOG_STD
|
|
|
|
type
|
|
MLP* = object
|
|
w1*, b1*: Tensor[float32] # [hidden, input], [hidden]
|
|
w2*, b2*: Tensor[float32] # [hidden, hidden], [hidden]
|
|
w3*, b3*: Tensor[float32] # [output, hidden], [output]
|
|
|
|
ActorCritic* = object
|
|
actor*: MLP
|
|
critic*: MLP
|
|
logStd*: Tensor[float32] # [ACTION_DIM] — one per action dim
|
|
|
|
proc initMLP*(inputDim, hiddenDim, outputDim: int): MLP =
|
|
# Xavier/He-style init: scale weights by sqrt(2/fan_in)
|
|
result.w1 = randomNormalTensor[float32]([hiddenDim, inputDim]) *. sqrt(2.0'f32 / inputDim.float32)
|
|
result.b1 = zeros[float32](hiddenDim)
|
|
result.w2 = randomNormalTensor[float32]([hiddenDim, hiddenDim]) *. sqrt(2.0'f32 / hiddenDim.float32)
|
|
result.b2 = zeros[float32](hiddenDim)
|
|
result.w3 = randomNormalTensor[float32]([outputDim, hiddenDim]) *. sqrt(1.0'f32 / hiddenDim.float32)
|
|
result.b3 = zeros[float32](outputDim)
|
|
|
|
proc initActorCritic*(): ActorCritic =
|
|
result.actor = initMLP(STATE_DIM, 64, ACTION_DIM)
|
|
result.critic = initMLP(STATE_DIM, 64, 1)
|
|
result.logStd = newTensor[float32](ACTION_DIM).map(proc(v: float32): float32 = initialLogStd)
|
|
|
|
proc forward*(mlp: MLP, x: Tensor[float32]): Tensor[float32] =
|
|
## x shape: [inputDim] (1D vector)
|
|
let h1 = tanh(mlp.w1 * x + mlp.b1)
|
|
let h2 = tanh(mlp.w2 * h1 + mlp.b2)
|
|
result = mlp.w3 * h2 + mlp.b3
|
|
|
|
proc actorForward*(ac: ActorCritic, state: Tensor[float32], deterministic = false): tuple[actions: Tensor[float32], logProb: float32] =
|
|
## state: [STATE_DIM]. Returns actions [ACTION_DIM] and sum log-prob.
|
|
## deterministic=true: return mean only (no noise), logProb=0.
|
|
let mean = ac.actor.forward(state)
|
|
if deterministic:
|
|
return (actions: mean, logProb: 0.0'f32)
|
|
|
|
# Floor logStd at logStdFloor before exp → min std ≈ exp(logStdFloor)
|
|
var clampedLogStd = ac.logStd.map(proc(v: float32): float32 = clamp(v, logStdFloor, logStdCeiling))
|
|
let std = clampedLogStd.map(proc(v: float32): float32 = exp(v))
|
|
|
|
var actions = newTensor[float32](ACTION_DIM)
|
|
var logP = 0.0'f32
|
|
for i in 0..<ACTION_DIM:
|
|
let mu = mean[i]
|
|
let s = std[i]
|
|
let z = gauss(0.0'f64, 1.0'f64).float32
|
|
actions[i] = mu + s * z
|
|
# log N(a; mu, s) = -0.5*((a-mu)/s)^2 - log(s) - 0.5*log(2π)
|
|
let diff = (actions[i] - mu) / s
|
|
logP += -0.5'f32 * diff * diff - ln(s) - 0.5'f32 * ln(2.0'f32 * PI.float32)
|
|
|
|
result = (actions: actions, logProb: logP)
|
|
|
|
proc criticForward*(ac: ActorCritic, state: Tensor[float32]): float32 =
|
|
## state: [STATE_DIM]. Returns scalar value estimate.
|
|
let val = ac.critic.forward(state)
|
|
result = val[0]
|
|
|
|
proc computeLogProb*(ac: ActorCritic, state, action: Tensor[float32]): float32 =
|
|
## Log-probability of action under current policy (no sampling).
|
|
let mean = ac.actor.forward(state)
|
|
let logStdClamped = ac.logStd.map(proc(v: float32): float32 = clamp(v, logStdFloor, logStdCeiling))
|
|
let std = logStdClamped.map(proc(v: float32): float32 = exp(v))
|
|
var logP = 0.0'f32
|
|
for i in 0..<ACTION_DIM:
|
|
let mu = mean[i]
|
|
let s = std[i]
|
|
let diff = (action[i] - mu) / s
|
|
logP += -0.5'f32 * diff * diff - ln(s) - 0.5'f32 * ln(2.0'f32 * PI.float32)
|
|
result = logP
|