Current section

Files

Jump to
macula_tweann src brain_learner.erl
Raw

src/brain_learner.erl

%% @doc Brain learner process for weight adaptation via plasticity.
%%
%% This GenServer manages the learning aspects of a brain system:
%% - Applies plasticity rules to update weights based on neural activity
%% - Maintains an experience buffer for batch learning
%% - Handles reward signals for reinforcement-style learning
%%
%% == Online Learning ==
%%
%% When online learning is enabled, the learner receives activation
%% data after each inference and applies plasticity rules:
%%
%% Inference → Activations → Learner → Weight Updates → Back to Inference
%%
%% == Batch Learning ==
%%
%% For delayed rewards (e.g., end of game), the learner buffers
%% experiences and applies learning when a reward is received:
%%
%% 1. Record experiences during episode
%% 2. Receive final reward
%% 3. Apply learning with reward propagation (eligibility traces)
%%
%% == Theory ==
%%
%% This module implements reward-modulated Hebbian learning, where
%% weight changes depend on:
%% - Pre-synaptic activity (input to connection)
%% - Post-synaptic activity (output from connection)
%% - Global reward signal (from environment)
%%
%% The basic rule: delta_w = learning_rate * pre * post * reward
%%
%% For delayed rewards, eligibility traces track which synapses
%% were recently active, allowing credit assignment across time.
%%
%% @see plasticity
%% @see plasticity_modulated
%% @copyright 2025 Macula.io
-module(brain_learner).
-behaviour(gen_server).
%% API
-export([
start_link/1,
stop/1,
%% Learning control
enable/1,
disable/1,
is_enabled/1,
%% Reward handling
reward/2,
set_baseline_reward/2,
%% Experience buffer
record_experience/3,
clear_experience/1,
learn_from_experience/1,
learn_from_experience/2,
get_experience_count/1,
set_auto_record/2,
get_auto_record/1,
%% Configuration
set_plasticity_rule/2,
get_plasticity_rule/1,
set_learning_rate/2,
get_learning_rate/1,
%% Weight access
get_weight_deltas/1
]).
%% gen_server callbacks
-export([
init/1,
handle_call/3,
handle_cast/2,
handle_info/2,
terminate/2
]).
-record(state, {
%% Brain ID for pubsub communication
brain_id :: term(),
%% Current network weights (updated via pubsub)
current_weights :: [float()] | undefined,
%% Learning configuration
enabled :: boolean(),
plasticity_rule :: atom(),
learning_rate :: float(),
baseline_reward :: float(),
%% Experience buffer for batch learning
experience_buffer :: [experience()],
max_buffer_size :: pos_integer(),
%% Current reward (for online learning)
current_reward :: float(),
%% Accumulated weight deltas (for inspection/debugging)
weight_deltas :: [float()],
%% Auto-record experiences from pubsub evaluated events
auto_record :: boolean()
}).
-type experience() :: #{
inputs := [float()],
activations := [[float()]],
outputs := [float()],
timestamp := integer()
}.
-export_type([experience/0]).
%%==============================================================================
%% API Functions
%%==============================================================================
%% @doc Start a brain learner process.
%%
%% Options:
%% - `inference_pid' - PID of the brain inference process (required for weight updates)
%% - `enabled' - Whether learning is enabled (default: true)
%% - `plasticity_rule' - Atom identifying the rule (default: modulated)
%% - `learning_rate' - Learning rate (default: 0.01)
%% - `baseline_reward' - Baseline to subtract from rewards (default: 0.0)
%% - `max_buffer_size' - Max experiences to buffer (default: 1000)
-spec start_link(map()) -> {ok, pid()} | {error, term()}.
start_link(Opts) ->
gen_server:start_link(?MODULE, Opts, []).
%% @doc Stop the learner process.
-spec stop(pid()) -> ok.
stop(Pid) ->
gen_server:stop(Pid).
%% @doc Enable learning.
-spec enable(pid()) -> ok.
enable(Pid) ->
gen_server:call(Pid, enable).
%% @doc Disable learning.
-spec disable(pid()) -> ok.
disable(Pid) ->
gen_server:call(Pid, disable).
%% @doc Check if learning is enabled.
-spec is_enabled(pid()) -> boolean().
is_enabled(Pid) ->
gen_server:call(Pid, is_enabled).
%% @doc Provide a reward signal.
%%
%% For online learning, this affects the next weight update.
%% Positive rewards strengthen active connections, negative weaken them.
%%
%% @param Pid Learner process
%% @param Reward Reward value (typically -1.0 to 1.0)
-spec reward(pid(), float()) -> ok.
reward(Pid, Reward) ->
gen_server:cast(Pid, {reward, Reward}).
%% @doc Set the baseline reward for comparison.
%%
%% Effective reward = actual_reward - baseline_reward.
%% This helps with reward normalization.
-spec set_baseline_reward(pid(), float()) -> ok.
set_baseline_reward(Pid, Baseline) ->
gen_server:call(Pid, {set_baseline_reward, Baseline}).
%% @doc Record an experience for batch learning.
%%
%% @param Pid Learner process
%% @param Inputs The sensor inputs
%% @param Activations All layer activations from the inference
-spec record_experience(pid(), [float()], [[float()]]) -> ok.
record_experience(Pid, Inputs, Activations) ->
gen_server:cast(Pid, {record_experience, Inputs, Activations}).
%% @doc Clear the experience buffer.
-spec clear_experience(pid()) -> ok.
clear_experience(Pid) ->
gen_server:call(Pid, clear_experience).
%% @doc Learn from buffered experiences using current reward.
-spec learn_from_experience(pid()) -> {ok, non_neg_integer()}.
learn_from_experience(Pid) ->
gen_server:call(Pid, learn_from_experience).
%% @doc Learn from buffered experiences with a specific reward.
%%
%% @param Pid Learner process
%% @param FinalReward Reward to apply (e.g., end-of-episode reward)
%% @returns {ok, NumExperiencesProcessed}
-spec learn_from_experience(pid(), float()) -> {ok, non_neg_integer()}.
learn_from_experience(Pid, FinalReward) ->
gen_server:call(Pid, {learn_from_experience, FinalReward}).
%% @doc Get the number of buffered experiences.
-spec get_experience_count(pid()) -> non_neg_integer().
get_experience_count(Pid) ->
gen_server:call(Pid, get_experience_count).
%% @doc Enable or disable automatic experience recording.
%%
%% When enabled, the learner automatically records experiences from
%% 'evaluated' events published by the brain via pubsub.
-spec set_auto_record(pid(), boolean()) -> ok.
set_auto_record(Pid, Enabled) ->
gen_server:call(Pid, {set_auto_record, Enabled}).
%% @doc Check if automatic experience recording is enabled.
-spec get_auto_record(pid()) -> boolean().
get_auto_record(Pid) ->
gen_server:call(Pid, get_auto_record).
%% @doc Set the plasticity rule.
%%
%% Available rules: none, hebbian, modulated
-spec set_plasticity_rule(pid(), atom()) -> ok.
set_plasticity_rule(Pid, Rule) ->
gen_server:call(Pid, {set_plasticity_rule, Rule}).
%% @doc Get the current plasticity rule.
-spec get_plasticity_rule(pid()) -> atom().
get_plasticity_rule(Pid) ->
gen_server:call(Pid, get_plasticity_rule).
%% @doc Set the learning rate.
-spec set_learning_rate(pid(), float()) -> ok.
set_learning_rate(Pid, Rate) ->
gen_server:call(Pid, {set_learning_rate, Rate}).
%% @doc Get the current learning rate.
-spec get_learning_rate(pid()) -> float().
get_learning_rate(Pid) ->
gen_server:call(Pid, get_learning_rate).
%% @doc Get accumulated weight deltas from last learning step.
%%
%% Useful for debugging and visualization.
-spec get_weight_deltas(pid()) -> [float()].
get_weight_deltas(Pid) ->
gen_server:call(Pid, get_weight_deltas).
%%==============================================================================
%% gen_server Callbacks
%%==============================================================================
init(Opts) ->
BrainId = maps:get(brain_id, Opts),
Enabled = maps:get(enabled, Opts, true),
PlasticityRule = maps:get(plasticity_rule, Opts, modulated),
LearningRate = maps:get(learning_rate, Opts, 0.01),
BaselineReward = maps:get(baseline_reward, Opts, 0.0),
MaxBufferSize = maps:get(max_buffer_size, Opts, 1000),
AutoRecord = maps:get(auto_record, Opts, true),
%% Subscribe to evaluated events from brain
brain_pubsub:subscribe(BrainId, evaluated),
%% Subscribe to reward events from brain_system
brain_pubsub:subscribe(BrainId, reward_received),
%% Request current weights from brain
brain_pubsub:subscribe(BrainId, weights_response),
brain_pubsub:publish(BrainId, weights_requested, #{}),
State = #state{
brain_id = BrainId,
current_weights = undefined,
enabled = Enabled,
plasticity_rule = PlasticityRule,
learning_rate = LearningRate,
baseline_reward = BaselineReward,
experience_buffer = [],
max_buffer_size = MaxBufferSize,
current_reward = 0.0,
weight_deltas = [],
auto_record = AutoRecord
},
{ok, State}.
handle_call(enable, _From, State) ->
{reply, ok, State#state{enabled = true}};
handle_call(disable, _From, State) ->
{reply, ok, State#state{enabled = false}};
handle_call(is_enabled, _From, State) ->
{reply, State#state.enabled, State};
handle_call({set_baseline_reward, Baseline}, _From, State) ->
{reply, ok, State#state{baseline_reward = Baseline}};
handle_call(clear_experience, _From, State) ->
{reply, ok, State#state{experience_buffer = []}};
handle_call(learn_from_experience, _From, State) ->
{NumProcessed, NewState} = do_learn_from_experience(State, State#state.current_reward),
{reply, {ok, NumProcessed}, NewState};
handle_call({learn_from_experience, FinalReward}, _From, State) ->
{NumProcessed, NewState} = do_learn_from_experience(State, FinalReward),
{reply, {ok, NumProcessed}, NewState};
handle_call(get_experience_count, _From, State) ->
{reply, length(State#state.experience_buffer), State};
handle_call({set_plasticity_rule, Rule}, _From, State) ->
{reply, ok, State#state{plasticity_rule = Rule}};
handle_call(get_plasticity_rule, _From, State) ->
{reply, State#state.plasticity_rule, State};
handle_call({set_learning_rate, Rate}, _From, State) ->
{reply, ok, State#state{learning_rate = Rate}};
handle_call(get_learning_rate, _From, State) ->
{reply, State#state.learning_rate, State};
handle_call(get_weight_deltas, _From, State) ->
{reply, State#state.weight_deltas, State};
handle_call({set_auto_record, Enabled}, _From, State) ->
{reply, ok, State#state{auto_record = Enabled}};
handle_call(get_auto_record, _From, State) ->
{reply, State#state.auto_record, State};
handle_call(_Request, _From, State) ->
{reply, {error, unknown_request}, State}.
handle_cast({reward, Reward}, State) ->
{noreply, State#state{current_reward = Reward}};
handle_cast({record_experience, Inputs, Activations}, State) ->
#state{
experience_buffer = Buffer,
max_buffer_size = MaxSize
} = State,
Experience = #{
inputs => Inputs,
activations => Activations,
outputs => lists:last(Activations),
timestamp => erlang:system_time(millisecond)
},
%% Add to buffer, trimming if necessary
NewBuffer = case length(Buffer) >= MaxSize of
true -> [Experience | lists:droplast(Buffer)];
false -> [Experience | Buffer]
end,
{noreply, State#state{experience_buffer = NewBuffer}};
handle_cast(_Msg, State) ->
{noreply, State}.
%% Handle evaluated events from brain (auto-record experiences)
handle_info({brain_event, evaluated, EventData}, State) ->
#state{
auto_record = AutoRecord,
experience_buffer = Buffer,
max_buffer_size = MaxSize
} = State,
NewState = case AutoRecord of
true ->
#{inputs := Inputs, activations := Activations, outputs := Outputs} = EventData,
Experience = #{
inputs => Inputs,
activations => Activations,
outputs => Outputs,
timestamp => erlang:system_time(millisecond)
},
NewBuffer = case length(Buffer) >= MaxSize of
true -> [Experience | lists:droplast(Buffer)];
false -> [Experience | Buffer]
end,
State#state{experience_buffer = NewBuffer};
false ->
State
end,
{noreply, NewState};
%% Handle reward events from brain_system
handle_info({brain_event, reward_received, #{reward := Reward}}, State) ->
{noreply, State#state{current_reward = Reward}};
%% Handle weights response (for initial sync)
handle_info({brain_event, weights_response, #{weights := Weights}}, State) ->
{noreply, State#state{current_weights = Weights}};
%% Ignore other pubsub events
handle_info({brain_event, _Topic, _Data}, State) ->
{noreply, State};
handle_info(_Info, State) ->
{noreply, State}.
terminate(_Reason, State) ->
%% Cleanup pubsub subscriptions
#state{brain_id = BrainId} = State,
brain_pubsub:cleanup(BrainId),
ok.
%%==============================================================================
%% Internal Functions
%%==============================================================================
%% @private Process buffered experiences and apply learning
do_learn_from_experience(State, _Reward) when State#state.enabled == false ->
{0, State};
do_learn_from_experience(State, _Reward) when State#state.experience_buffer == [] ->
{0, State};
do_learn_from_experience(State, _Reward) when State#state.current_weights == undefined ->
%% No weights available yet, can't learn
{0, State};
do_learn_from_experience(State, Reward) ->
#state{
brain_id = BrainId,
current_weights = Weights,
experience_buffer = Buffer,
plasticity_rule = Rule,
learning_rate = LearningRate,
baseline_reward = Baseline
} = State,
%% Calculate effective reward
EffectiveReward = Reward - Baseline,
%% Convert flat weights to weight_spec format for plasticity
WeightSpecs = weights_to_specs(Weights, LearningRate),
%% Apply learning from each experience
{UpdatedSpecs, Deltas} = apply_learning_to_experiences(
WeightSpecs,
lists:reverse(Buffer), %% Oldest first
Rule,
EffectiveReward
),
%% Convert back to flat weights
NewWeights = specs_to_weights(UpdatedSpecs),
%% Publish weights_updated event for brain to apply
brain_pubsub:publish(BrainId, weights_updated, #{
weights => NewWeights,
deltas => Deltas
}),
NumProcessed = length(Buffer),
NewState = State#state{
experience_buffer = [],
current_weights = NewWeights, %% Update our copy too
weight_deltas = Deltas
},
{NumProcessed, NewState}.
%% @private Convert flat weights to weight_spec tuples
weights_to_specs(Weights, LearningRate) ->
[{W, 0.0, LearningRate, []} || W <- Weights].
%% @private Convert weight_spec tuples back to flat weights
specs_to_weights(Specs) ->
[plasticity:get_weight(Spec) || Spec <- Specs].
%% @private Apply learning across all experiences
apply_learning_to_experiences(WeightSpecs, [], _Rule, _Reward) ->
Deltas = [plasticity:get_delta(Spec) || Spec <- WeightSpecs],
{WeightSpecs, Deltas};
apply_learning_to_experiences(WeightSpecs, [Experience | Rest], Rule, Reward) ->
#{activations := Activations} = Experience,
%% Apply plasticity rule to each weight based on activations
UpdatedSpecs = apply_plasticity_to_weights(WeightSpecs, Activations, Rule, Reward),
apply_learning_to_experiences(UpdatedSpecs, Rest, Rule, Reward).
%% @private Apply plasticity rule to weights based on layer activations
%%
%% This maps the flat weight list to the layer structure and applies
%% the plasticity rule based on pre/post activations.
apply_plasticity_to_weights(WeightSpecs, Activations, Rule, Reward) ->
RuleModule = plasticity:rule_module(Rule),
%% For a feedforward network, weights connect adjacent layers
%% We need to pair pre-activations (layer N) with post-activations (layer N+1)
NumLayers = length(Activations),
case NumLayers < 2 of
true ->
%% Not enough layers for learning
WeightSpecs;
false ->
%% Build list of {PreActivations, PostActivations} pairs
LayerPairs = lists:zip(
lists:droplast(Activations), %% Pre: all but last
tl(Activations) %% Post: all but first
),
%% Apply plasticity to each weight
%% Weights are organized: [layer1_weights..., layer2_weights..., ...]
%% Each layer has: [w11, w12, ..., w1n, w21, ..., bias1, bias2, ...]
apply_plasticity_by_layer(WeightSpecs, LayerPairs, RuleModule, Reward)
end.
%% @private Apply plasticity layer by layer
apply_plasticity_by_layer(WeightSpecs, [], _RuleModule, _Reward) ->
WeightSpecs;
apply_plasticity_by_layer(WeightSpecs, [{PreActs, PostActs} | RestPairs], RuleModule, Reward) ->
PreSize = length(PreActs),
PostSize = length(PostActs),
%% Number of weights in this layer (excluding biases)
WeightCount = PreSize * PostSize,
BiasCount = PostSize,
TotalCount = WeightCount + BiasCount,
%% Split weights for this layer
{LayerSpecs, RestSpecs} = lists:split(TotalCount, WeightSpecs),
{ConnectionSpecs, BiasSpecs} = lists:split(WeightCount, LayerSpecs),
%% Apply plasticity to connection weights
%% Weights are stored as: [w(1,1), w(1,2), ..., w(1,post), w(2,1), ...]
%% where w(pre_idx, post_idx)
UpdatedConnections = apply_plasticity_to_connections(
ConnectionSpecs, PreActs, PostActs, RuleModule, Reward
),
%% Biases: use post-activation only (no pre-synaptic activity)
UpdatedBiases = lists:zipwith(
fun(BiasSpec, PostAct) ->
RuleModule:apply_rule(BiasSpec, 1.0, PostAct, Reward)
end,
BiasSpecs,
PostActs
),
%% Combine and continue with next layer
UpdatedLayerSpecs = UpdatedConnections ++ UpdatedBiases,
apply_plasticity_by_layer(UpdatedLayerSpecs ++ RestSpecs, RestPairs, RuleModule, Reward).
%% @private Apply plasticity to connection weights
apply_plasticity_to_connections(ConnectionSpecs, PreActs, PostActs, RuleModule, Reward) ->
PreSize = length(PreActs),
PostSize = length(PostActs),
%% Map over all connections
lists:zipwith(
fun(Spec, Idx) ->
%% Calculate which pre/post neurons this weight connects
PreIdx = ((Idx - 1) rem PreSize) + 1,
PostIdx = ((Idx - 1) div PreSize) + 1,
PreAct = lists:nth(PreIdx, PreActs),
PostAct = lists:nth(PostIdx, PostActs),
RuleModule:apply_rule(Spec, PreAct, PostAct, Reward)
end,
ConnectionSpecs,
lists:seq(1, PreSize * PostSize)
).