Documentation
¶
Index ¶
- Constants
- Variables
- func NoisyLinear(ctx *context.Context, x *Node, outFeatures int, sigmaInit float64) *Node
- func NoisyLinearEval(ctx *context.Context, x *Node, outFeatures int, sigmaInit float64) *Node
- func NoisyLinearExternal(ctx *context.Context, x, epsIn, epsOut *Node, outFeatures int, ...) *Node
- type BetaSchedule
- type PrioritizedReplayBuffer
- func (b *PrioritizedReplayBuffer) Cap() int
- func (b *PrioritizedReplayBuffer) Len() int
- func (b *PrioritizedReplayBuffer) LoadFromBuffer(buf []Transition)
- func (b *PrioritizedReplayBuffer) Push(t Transition)
- func (b *PrioritizedReplayBuffer) Sample(batchSize int, beta float64) ([]Transition, []int, []float64)
- func (b *PrioritizedReplayBuffer) ToBuffer() []Transition
- func (b *PrioritizedReplayBuffer) UpdatePriority(idx int, tdError float64)
- type ReplayBuffer
- type RunningNormalizer
- type SumTree
- type TargetNetwork
- type TargetNetworkOption
- type Transition
Constants ¶
const DefaultNoisySigmaInit = 0.5
Variables ¶
return backends.MustNew() })
SharedBackend returns the process-global XLA backend, creating it on first call. All DQN balancers, schedulers, and TD3 optimizers share this single backend to avoid native-memory leaks from per-instance churn (PJRT CPU client + device buffer arena + compiled program cache are expensive to allocate/destroy).
Functions ¶
func NoisyLinear ¶
NoisyLinear implements a Factorised Gaussian Noisy Linear layer (Fortunato et al., ICLR 2018).
Instead of external epsilon-greedy or softmax exploration, the noise is baked into the network weights themselves. Each forward pass samples fresh noise via RandomNormal, and the learnable sigma parameters control the noise amplitude per-weight. As training progresses and the network becomes more confident, sigma shrinks toward zero, making the policy naturally greedy.
The factorised variant reduces the noise parameter count from (p*q + q) to (p + q) by decomposing the noise matrix as:
f(eps_in) ⊗ f(eps_out) where f(x) = sign(x) * sqrt(|x|)
func NoisyLinearEval ¶
NoisyLinearEval performs inference using only the mu (mean) weights, without sampling noise. Useful for final evaluation metrics or ONNX export after training is complete. NOT suitable as the main inference path during training — see NoisyLinearExternal for exploration-preserving inference.
func NoisyLinearExternal ¶
func NoisyLinearExternal(ctx *context.Context, x, epsIn, epsOut *Node, outFeatures int, sigmaInit float64) *Node
NoisyLinearExternal is identical to NoisyLinear except that epsIn and epsOut are provided as graph *Node inputs sampled by the caller (e.g. via math/rand/v2.NormFloat64 on the CPU) rather than ctx.RandomNormal.
This avoids the gomlx #rngState concurrency issue (concurrent Exec objects that share a context call Reset() on the RNG variable during setSideParams, racing with each other) while preserving exploration noise on every forward pass — which is the core mechanism of NoisyNets (Fortunato 2017 §4).
Types ¶
type BetaSchedule ¶
type BetaSchedule struct {
// contains filtered or unexported fields
}
BetaSchedule linearly anneals beta from start to end over a given number of steps.
func NewBetaSchedule ¶
func NewBetaSchedule(start, end float64, steps int) *BetaSchedule
func (*BetaSchedule) Next ¶
func (b *BetaSchedule) Next() float64
type PrioritizedReplayBuffer ¶
type PrioritizedReplayBuffer struct {
// contains filtered or unexported fields
}
PrioritizedReplayBuffer samples transitions proportional to priority^alpha using a SumTree for O(log N) sampling (vs O(N) linear scan).
Not safe for concurrent use; designed for the single-writer trainLoop pattern where Push and Sample are called from the same goroutine.
func NewPrioritizedReplayBuffer ¶
func NewPrioritizedReplayBuffer(capacity int, alpha float64) *PrioritizedReplayBuffer
func (*PrioritizedReplayBuffer) Cap ¶
func (b *PrioritizedReplayBuffer) Cap() int
func (*PrioritizedReplayBuffer) Len ¶
func (b *PrioritizedReplayBuffer) Len() int
Len is safe to call from any goroutine.
func (*PrioritizedReplayBuffer) LoadFromBuffer ¶
func (b *PrioritizedReplayBuffer) LoadFromBuffer(buf []Transition)
LoadFromBuffer imports previously cached transitions.
func (*PrioritizedReplayBuffer) Push ¶
func (b *PrioritizedReplayBuffer) Push(t Transition)
func (*PrioritizedReplayBuffer) Sample ¶
func (b *PrioritizedReplayBuffer) Sample(batchSize int, beta float64) ([]Transition, []int, []float64)
Sample returns a batch of transitions sampled proportional to priority^alpha using stratified sampling over the SumTree. O(B * log N) total. IS-weights are computed as w_i = (N * P(i))^(-beta) / max(w), normalized so max=1.0.
func (*PrioritizedReplayBuffer) ToBuffer ¶
func (b *PrioritizedReplayBuffer) ToBuffer() []Transition
ToBuffer exports all transitions for warmup caching.
func (*PrioritizedReplayBuffer) UpdatePriority ¶
func (b *PrioritizedReplayBuffer) UpdatePriority(idx int, tdError float64)
type ReplayBuffer ¶
type ReplayBuffer struct {
// contains filtered or unexported fields
}
func NewReplayBuffer ¶
func NewReplayBuffer(capacity int) *ReplayBuffer
func (*ReplayBuffer) Len ¶
func (rb *ReplayBuffer) Len() int
func (*ReplayBuffer) Push ¶
func (rb *ReplayBuffer) Push(t Transition)
func (*ReplayBuffer) Sample ¶
func (rb *ReplayBuffer) Sample(batchSize int) []Transition
Sample returns up to batchSize random transitions.
func (*ReplayBuffer) SampleInto ¶
func (rb *ReplayBuffer) SampleInto(dst []Transition) int
type RunningNormalizer ¶
type RunningNormalizer struct {
// contains filtered or unexported fields
}
RunningNormalizer computes online mean and variance using Welford's algorithm and normalizes vectors to zero mean and unit variance.
Observe and Reset must be called from a single goroutine (the policy loop). NormalizeInPlace and NormalizeBatch may be called from any goroutine.
func NewRunningNormalizer ¶
func NewRunningNormalizer(size int) *RunningNormalizer
func (*RunningNormalizer) Normalize ¶
func (rn *RunningNormalizer) Normalize(x []float64) []float64
Normalize returns (x - mean) / (std + ε) without modifying x.
func (*RunningNormalizer) NormalizeBatch ¶
func (rn *RunningNormalizer) NormalizeBatch(buf []float64, batchSize int)
NormalizeBatch normalizes a flat buffer of concatenated state vectors in-place.
func (*RunningNormalizer) NormalizeInPlace ¶
func (rn *RunningNormalizer) NormalizeInPlace(x []float64)
NormalizeInPlace normalizes x in-place: x[i] = (x[i] - mean[i]) / (std[i] + ε).
func (*RunningNormalizer) Observe ¶
func (rn *RunningNormalizer) Observe(x []float64)
Must be called from the single-writer goroutine only.
func (*RunningNormalizer) Reset ¶
func (rn *RunningNormalizer) Reset()
Must be called from the single-writer goroutine only.
type SumTree ¶
type SumTree struct {
// contains filtered or unexported fields
}
SumTree is a segment tree that supports O(log N) priority-weighted sampling and O(log N) priority updates (Schaul et al., ICML 2016).
func NewSumTree ¶
type TargetNetwork ¶
type TargetNetwork struct {
// contains filtered or unexported fields
}
TargetNetwork maintains a separate set of weights for computing stable target Q-values (Mnih 2015). It periodically performs a hard copy from the main context.
func NewTargetNetwork ¶
func NewTargetNetwork(mainCtx *context.Context, opts ...TargetNetworkOption) *TargetNetwork
NewTargetNetwork creates a target network by cloning all trainable variables from mainCtx. Must be called after mainCtx variables are initialized (i.e., after a dummy forward pass).
func (*TargetNetwork) HardSync ¶
func (tn *TargetNetwork) HardSync()
func (*TargetNetwork) Step ¶
func (tn *TargetNetwork) Step()
Step increments the internal counter and performs a hard sync every syncEvery steps.
func (*TargetNetwork) TargetCtx ¶
func (tn *TargetNetwork) TargetCtx() *context.Context
type TargetNetworkOption ¶
type TargetNetworkOption func(*TargetNetwork)
func WithSyncEvery ¶
func WithSyncEvery(n int) TargetNetworkOption
Source Files