Study Plan

Apply One Policy-Value Head Training Step

Reinforcement LearningNeural NetworksOptimization

Hard

Apply one SGD step to policy and value heads on fixed features. Let $X$ be features of shape $(B,D)$ , where $B$ is batch size and $D$ is feature width. With $A$ actions, policy_weight $W_p$ has shape $(A,D)$ and policy_bias $b_p$ has shape $(A,)$ . value_weight $w_v$ has shape $(D,)$ ; value_bias $b_v$ is a scalar tensor or has shape $(1,)$ . The heads are

p=\operatorname{softmax}(XW_p^{\mathsf T}+b_p).

v=\tanh(Xw_v+b_v).

Here $p$ contains action probabilities, with softmax across actions, and $v$ contains predicted values. For target_policies $\pi$ of shape $(B,A)$ and target_values $z$ of shape $(B,)$ , the loss is

\ell=-\frac{1}{B}\sum_{i=1}^{B}\sum_{a=1}^{A}\pi_{ia}\log p_{ia}+\frac{1}{B}\sum_{i=1}^{B}(v_i-z_i)^2+\lambda\sum_{\theta\in\Theta}\lVert\theta\rVert_2^2.

Here $i$ indexes examples, $a$ indexes actions, $\lambda$ is l2_coefficient, and $\Theta$ contains both weights and both biases; each squared norm sums all squared entries. Update all four parameters from the same original loss:

\theta'=\theta-\eta\nabla_{\theta}\ell.

Here $\eta$ is learning_rate and $\theta'$ is an updated parameter. Features and targets remain fixed. Return a five-item tuple of detached PyTorch float32 tensors: the scalar pre-update loss, updated policy_weight, policy_bias, value_weight, and value_bias. Preserve each parameter's shape and the input device.

Loading visualization...

Examples

Input: features=[[1,0]], policy_weight=[[0.2,-0.1],[-0.3,0.4]], policy_bias=[0.1,-0.2], value_weight=[0.5,-0.5], value_bias=[0], target_policies=[[1,0]], target_values=[1], learning_rate=0.1, l2_coefficient=0.01

Output: (0.668918, [[0.230603,-0.0998],[-0.330403,0.3992]], [0.130803,-0.230603], [0.583603,-0.499], [0.084603])

Explanation: The returned loss is measured before updating either head; both biases are included in the regularization.

Input: features=[[0,1],[1,1]], policy_weight=[[0,0],[0,0]], policy_bias=[0,0], value_weight=[0,0], value_bias=[0], target_policies=[[0.25,0.75],[0.5,0.5]], target_values=[-1,0.5], learning_rate=0, l2_coefficient=0

Output: (1.318147, [[0,0],[0,0]], [0,0], [0,0], [0])

Input: features=[[2]], policy_weight=[[0.5],[-0.5],[0.25]], policy_bias=[0,0,0], value_weight=[0.1], value_bias=0, target_policies=[[0.2,0.3,0.5]], target_values=[0], learning_rate=0.05, l2_coefficient=0.2

Output: (1.558414, [[0.45259],[-0.46777],[0.260179]], [-0.018705,0.011115,0.00759], [0.060063], -0.018969)

Hint 1

torch.log_softmax gives stable log probabilities across actions.

Hint 2

torch.autograd.grad can return gradients for a tuple of parameters.

Hint 3

Tensor.detach removes the returned tensor from the gradient graph.

Requirements

Use PyTorch autograd and float32 calculations.
Include both weights and both biases in regularization.
Leave every input unchanged and return tensors detached from the gradient graph.

Constraints

All tensor inputs are finite float32 tensors on the same device.
B, D, and A are positive; tensor shapes follow the description.
Each target-policy row is nonnegative and sums to one.
Target values lie between -1 and 1.
learning_rate and l2_coefficient are finite and nonnegative.

Try Similar Problems

Assemble a Compact Dual-Head NetworkHard Run Deterministic Neural-Guided MCTSHard Compute the Policy and Value HeadsMedium Apply an AlphaGo Residual BlockMedium Train a Tiny Micrograd MLPHard Apply a Dihedral Go SymmetryMedium Gradient Clipping (Global Norm)Medium Sample a Recent Self-Play Replay BatchEasy GRPO Group-Relative AdvantageHard Dropout from ScratchHard PPO with Reference-Policy KLHard Bradley-Terry Reward Model LossMedium

Apply One Policy-Value Head Training Step

Reinforcement LearningNeural NetworksOptimization

Hard

p=\operatorname{softmax}(XW_p^{\mathsf T}+b_p).

v=\tanh(Xw_v+b_v).

Here $p$ contains action probabilities, with softmax across actions, and $v$ contains predicted values. For target_policies $\pi$ of shape $(B,A)$ and target_values $z$ of shape $(B,)$ , the loss is

\ell=-\frac{1}{B}\sum_{i=1}^{B}\sum_{a=1}^{A}\pi_{ia}\log p_{ia}+\frac{1}{B}\sum_{i=1}^{B}(v_i-z_i)^2+\lambda\sum_{\theta\in\Theta}\lVert\theta\rVert_2^2.

\theta'=\theta-\eta\nabla_{\theta}\ell.

Loading visualization...

Examples

Output: (0.668918, [[0.230603,-0.0998],[-0.330403,0.3992]], [0.130803,-0.230603], [0.583603,-0.499], [0.084603])

Explanation: The returned loss is measured before updating either head; both biases are included in the regularization.

Output: (1.318147, [[0,0],[0,0]], [0,0], [0,0], [0])

Output: (1.558414, [[0.45259],[-0.46777],[0.260179]], [-0.018705,0.011115,0.00759], [0.060063], -0.018969)

Hint 1

torch.log_softmax gives stable log probabilities across actions.

Hint 2

torch.autograd.grad can return gradients for a tuple of parameters.

Hint 3

Tensor.detach removes the returned tensor from the gradient graph.

Requirements

Use PyTorch autograd and float32 calculations.
Include both weights and both biases in regularization.
Leave every input unchanged and return tensors detached from the gradient graph.

Constraints

All tensor inputs are finite float32 tensors on the same device.
B, D, and A are positive; tensor shapes follow the description.
Each target-policy row is nonnegative and sums to one.
Target values lie between -1 and 1.
learning_rate and l2_coefficient are finite and nonnegative.

Try Similar Problems

features =

Accepts: array

policy_weight =

Accepts: array

policy_bias =

Accepts: array

value_weight =

Accepts: array

value_bias =

Accepts: array, number

target_policies =

Accepts: array

target_values =

Accepts: array

learning_rate =

Accepts: number

l2_coefficient =

Accepts: number