Writing a custom loss function in pysr to find recursive, hidden function

[isavdh]

Hi,
I'm trying to write a custom loss function for the following problem:
I modeled the following dataset: variables [a, r] (action and reward). This datasets represents a multi-armed bandit task in which a subject is asked to choose one out of four arms that yield different results from different normal distributions. Their action is the chosen action and their result is the reward. I want to find the subjects' expectation q of the arms and how the expectation is updated based on their last reward. The functions I used to model the dataset look like this: 1) Pt(a_k) = {(1-epsilon) if k == argmax(Q), epsilon/3 else}, 2)q_{t+1,k} = q_{t,k} + alpha * (r_t - q_{t,k}). This last function I want to find in the dataset using symbolic regression in Pysr. However, since the variables I'm using and predicting are not directly in the dataset, this problem comes with some difficulties. To solve the problem I'm trying to write a custom loss function, but it's not working as I want, since the model is producing only solutions with a complexity of 1. The code I wrote looks like this:

X = df2[['a','r']].to_numpy()
Y = df2['a'].to_numpy()

my_loss_function = """
function my_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    X = dataset.X 
    y = dataset.y
    n_trials = size(X,1)
    
    Q = zeros(4)
    total_loss = zero(L)

    for i in 1:n_trials 
        a=Int(X[1,i]) +1  
        r=X[2,i]

        q_input = [Q[a], r]
        q_update, valid = eval_tree_array(tree, reshape(q_input, :, 1), options)
    
        if !valid || any(isnan, q_update) || any(isinf, q_update) 
            return L(Inf)  
        end

        maximum = argmax(Q)
        list_ = Float64[]

        for i in 1:4
            if i == maximum
                push!(list_, 1 - 0.7)
            else
                push!(list_, 0.7 / 3)
            end
        end

        total_loss += -log(list_[a])
        
        Q[a] = q_update[1]

    end

    return total_loss/n_trials 
end
"""

model = PySRRegressor(
    niterations=1000,
    binary_operators=["+", "-", "*", "/"],
    unary_operators=['cos','sin','log','abs'], 
    loss_function=my_loss_function,
)

model.fit(X, Y, variable_names=['a','r'])

I'm not sure why the model only produces solutions with a complexity of 1. I thought it might have something to do with the model being fit to data that is different than the variables I'm using in the eval_tree_array() function, but I don't know any other way to fix it. Hopefully someone can help me with this. Thanks in advance!

[MilesCranmer]

I wonder if it might be easier to solve this with a TemplateExpressionSpec as opposed to a custom loss function? https://ai.damtp.cam.ac.uk/pysr/examples/#11-expression-specifications Edit: nvm, this is probably not the right way.

[MilesCranmer]

Ah. Is the issue because of

n_trials = size(X, 1)

by any chance? I think this should be size(X, 2), right?

[MilesCranmer]

Some other suggested improvements:

function my_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
    X = dataset.X 
    n_trials = size(X, 2)  # <-- WITH FIX APPLIED
    Q = ntuple(j -> zero(T), 4)  # <-- TUPLES FASTER THAN VECTORS
    total_loss = zero(L)

    for i in 1:n_trials 
        a = Int(X[1, i]) +1  
        r = X[2, i]

        q_input = [Q[a]; r;;]  # <-- CONSTRUCT 2x1 DIRECTLY
        q_update, valid = eval_tree_array(tree, q_input, options)
    
        if !valid  # <-- OTHER CHECKS NOT NEEDED
            return L(Inf)  
        end

        maximum = argmax(Q)
        total_loss += a == maximum ? -log(1 - 0.7) : -log(0.7 / 3)  # <-- SIMPLIFIED
        
        Q = ntuple(
            let a=a, q_update=q_update, Q=Q  # (not 100% needed; but this is a performance hack to avoid boxing in closure)
                j -> j == a ? q_update[1] : Q[j]
            end,
            4
         )  # <-- CONSTRUCT NEW TUPLE
    end

    return total_loss / n_trials 
end

[isavdh]

Thank you so much for your answer, you really helped me out a lot and this code works perfectly! I'm quite new to programming in julia and using Pysr and I'm still left with one question: The functions this algorithm is producing contain the variable names 'a' and 'r', which are the names of the variables in X. However, I'm feeding eval_tree_array with different variables and I would like those variables(q[a] and r) to be used in the functions. I'm assuming the model is using the correct variables but assigning the wrong names to them. Is this a correct assumption? And how can I check which names the model is assigning to which variables? Thanks in advance for your answer!

[MilesCranmer]

I'm assuming the model is using the correct variables but assigning the wrong names to them. Is this a correct assumption?

Yes! The expression is represented using integers corresponding to inputs. So 1 would be the first input feature, 2 the second input feature, etc. The printing, for convenience, just assumes that 1 is the same as the first variable name of X. But if you are using a custom loss function this may not be true. But that's just the printing itself so don't worry about it