How to make smooth_l1_loss?

[wjessup]

Just got into this today but not much can be found searching and chatgpt isn't helpful.

Anyway, here's some cartpole code I'm converting from pytorch.

            # THIS STATE ACTION VALUES
            policy_this_preds  = policy_net(Tensor(state_batch))
            state_q_values = policy_this_preds.gather(Tensor(action_batch).unsqueeze(1), 1) 

            # NEXT STATE MAX VALUES
            Tensor.no_grad = True
            target_next_preds = target_net(Tensor(observation_batch)).max(axis=-1).numpy()
            Tensor.no_grad = False
            
            target_next_preds[terminated_batch] = 0
            expected_q_values = (target_next_preds * 0.99) + 1

            # smooth l1 loss
            diff = (state_q_values - Tensor(expected_q_values)).abs()
            less_than_delta = diff < 1.0
            more_than_delta = diff > 1.0
            diff = diff.where(less_than_delta, 0.5 * diff ** 2)
            loss = diff.where(more_than_delta, diff - 0.5 * 1.0)
            loss = loss.mean()

I get an error: RuntimeError: backward not implemented for <class 'tinygrad.mlops.Less'>

Also, is there a better way to do target_next_preds[terminated_batch] = 0 without needing to convert to numpy and back?

[chenyuxyz]

it's condition.where(value_if_true, value_if_false)

[wjessup]

updated to:

            diff = (state_q_values - Tensor(expected_q_values)).abs()
            loss = (diff < 1.0).where(0.5 * diff ** 2, diff - 0.5 * 1.0 )
            loss = loss.mean()

now getting: in backward assert (t0.grad is not None)

Due to converting expected_q_values to numpy?

[wjessup]

solved:

diff = Tensor.where(diff.detach() < 1.0, 0.5 * diff ** 2, diff - 0.5 * 1.0 ).mean()
```