Study Plan

Assemble a Compact Dual-Head Network

Reinforcement LearningNeural Networks

Hard

Build a compact AlphaGo Zero network from supplied parameter tensors. Let $B$ be the batch size, $N$ the board size, $C$ the trunk width, $C_p$ the policy-head width, $C_v$ the value-head width, and $H$ the value hidden width. The input X has shape (B, Cin, N, N), where Cin is its channel count.

The trunk uses trunk_weight for its 3 by 3 convolution and trunk_scale and trunk_shift for normalization, followed by ReLU. All convolutions use stride 1 and no bias; padding is 1 for 3 by 3 kernels and 0 for 1 by 1 kernels.

H_0=\operatorname{ReLU}\left(\operatorname{BN}\left(\operatorname{Conv}_{3\times3}(X)\right)\right)

For a feature tensor $Z$ , channel normalization computes the mean $\mu$ and population variance $\sigma^2$ over the batch and spatial dimensions. The supplied per-channel scale $\gamma$ , shift $\delta$ , and positive constant $\varepsilon$ are then applied as follows.

\operatorname{BN}(Z)=\gamma\frac{Z-\mu}{\sqrt{\sigma^2+\varepsilon}}+\delta

Each residual block applies two padded 3 by 3 convolutions. The first normalization is followed by ReLU, while the second is followed by the shortcut addition and ReLU. Every residual-block dictionary contains:

First convolution weights: conv1_weight
First normalization scale and shift: scale1 and shift1
Second convolution weights: conv2_weight
Second normalization scale and shift: scale2 and shift2

R_k=\operatorname{BN}_2\left(\operatorname{Conv}_{3\times3}\left(\operatorname{ReLU}\left(\operatorname{BN}_1\left(\operatorname{Conv}_{3\times3}(H_{k-1})\right)\right)\right)\right)

H_k=\operatorname{ReLU}(H_{k-1}+R_k)

Here $R_k$ is the residual branch output and $H_k$ is the shared feature tensor after block $k$ . The policy head uses a 1 by 1 convolution, normalization, ReLU, flattening, and a linear projection to produce one raw logit per board action plus pass. The value head uses its own 1 by 1 convolution, normalization, ReLU, hidden projection followed by ReLU, and tanh-bounded scalar projection. Flattening preserves the batch dimension and orders entries by channel, row, then column. Linear weights store output features by input features. The output-head dictionary contains:

Policy convolution weights: policy_conv_weight
Policy normalization scale and shift: policy_scale and policy_shift
Policy linear weights and bias: policy_linear_weight and policy_linear_bias
Value convolution weights: value_conv_weight
Value normalization scale and shift: value_scale and value_shift
Value hidden-layer weights and bias: value_hidden_weight and value_hidden_bias
Value output weights and bias: value_output_weight and value_output_bias

(L,v)=\operatorname{Heads}(H_K)

Here, K is the number of supplied residual blocks, L is the policy-logit tensor, and v is the value tensor. Return the tuple $(L,v)$ as 32-bit floating-point tensors on the input device, with shapes (B, N squared plus one) and (B,) respectively.

Loading visualization...

Examples

Input: x = [[[[1]]]], trunk_weight = [[[[0,0,0],[0,0,0],[0,0,0]]]], trunk_scale = [1], trunk_shift = [1], residual_parameters = [{"scale1":[1],"scale2":[1],"shift1":[0],"shift2":[0],"conv1_weight":[[[[0,0,0],[0,0,0],[0,0,0]]]],"conv2_weight":[[[[0,0,0],[0,0,0],[0,0,0]]]]}], head_parameters = {"value_scale":[1],"value_shift":[1],"policy_scale":[1],"policy_shift":[1],"value_conv_weight":[[[[0]]]],"value_hidden_bias":[0],"value_output_bias":[0],"policy_conv_weight":[[[[0]]]],"policy_linear_bias":[0,0],"value_hidden_weight":[[1]],"value_output_weight":[[1]],"policy_linear_weight":[[1],[2]]}, epsilon = 0.00001

Output: ([[1.0,2.0]], [0.761594])

Explanation: Each zero convolution is replaced by its normalization shift, and the value branch returns tanh(1).

Input: x shape = (1,1,1,1), trunk_weight shape = (1,1,3,3), trunk_scale = [1], trunk_shift = [0], residual_parameters = [], head_parameters contains policy_linear_weight shape (2,1) and value_hidden_weight shape (1,1), epsilon = 0.00001

Output: ([[0.5,-0.5]], [0.964028])

Input: x shape = (1,1,2,2), trunk_weight shape = (1,1,3,3), trunk_scale = [1], trunk_shift = [1], residual_parameters contains two blocks, head_parameters contains policy_linear_weight shape (5,4) and value_hidden_weight shape (1,4), epsilon = 0.00001

Output: (float32 tensor with shape (1,5), float32 tensor with shape (1,))

Hint 1

A normalization helper can be reused with each stage's own scale and shift.

Hint 2

Keep the block input available for the residual addition.

Hint 3

flatten(1) preserves one policy vector and one value per batch item.

Requirements

Use PyTorch float32 computation and outputs.
Residual blocks follow their supplied order; the list may be empty.
Return raw policy logits without softmax or legal-action masking.
Preserve all input tensors and their device.

Constraints

The input has shape $(B,C_{in},N,N)$ .
The trunk convolution weights have shape $(C,C_{in},3,3)$ , and its channel scale and shift each have shape $(C)$ .
Every residual convolution has weight shape $(C,C,3,3)$ , with channel scale and shift vectors of shape $(C)$ .
The policy convolution has weight shape $(C_p,C,1,1)$ , and its linear layer maps $C_pN^2$ values to $N^2+1$ logits.
The value convolution has weight shape $(C_v,C,1,1)$ , its hidden layer maps $C_vN^2$ values to $H$ , and its output layer maps $H$ values to one scalar.
The residual block list may be empty.
Every dimension is positive. All supplied tensor values are finite floating-point values on the same device; input dtypes may differ.
Residual parameters is a list of dictionaries and head parameters is a dictionary, with tensor values under the named keys.
Epsilon is a positive finite number.

Try Similar Problems

Compute the Policy and Value HeadsMedium Apply an AlphaGo Residual BlockMedium Apply One Policy-Value Head Training StepHard Encode AlphaGo Zero Board HistoryMedium Apply a Go Action with Captures and RepetitionHard Build the Legal Go Action MaskHard Run Deterministic Neural-Guided MCTSHard Turn Root Visits into a Search PolicyMedium DQN Loss with Target NetworkMedium Double DQN TargetEasy Dueling DQN DecompositionEasy GRPO Group-Relative AdvantageHard

Assemble a Compact Dual-Head Network

Reinforcement LearningNeural Networks

Hard

H_0=\operatorname{ReLU}\left(\operatorname{BN}\left(\operatorname{Conv}_{3\times3}(X)\right)\right)

\operatorname{BN}(Z)=\gamma\frac{Z-\mu}{\sqrt{\sigma^2+\varepsilon}}+\delta

First convolution weights: conv1_weight
First normalization scale and shift: scale1 and shift1
Second convolution weights: conv2_weight
Second normalization scale and shift: scale2 and shift2

R_k=\operatorname{BN}_2\left(\operatorname{Conv}_{3\times3}\left(\operatorname{ReLU}\left(\operatorname{BN}_1\left(\operatorname{Conv}_{3\times3}(H_{k-1})\right)\right)\right)\right)

H_k=\operatorname{ReLU}(H_{k-1}+R_k)

Policy convolution weights: policy_conv_weight
Policy normalization scale and shift: policy_scale and policy_shift
Policy linear weights and bias: policy_linear_weight and policy_linear_bias
Value convolution weights: value_conv_weight
Value normalization scale and shift: value_scale and value_shift
Value hidden-layer weights and bias: value_hidden_weight and value_hidden_bias
Value output weights and bias: value_output_weight and value_output_bias

(L,v)=\operatorname{Heads}(H_K)

Loading visualization...

Examples

Output: ([[1.0,2.0]], [0.761594])

Explanation: Each zero convolution is replaced by its normalization shift, and the value branch returns tanh(1).

Output: ([[0.5,-0.5]], [0.964028])

Output: (float32 tensor with shape (1,5), float32 tensor with shape (1,))

Hint 1

A normalization helper can be reused with each stage's own scale and shift.

Hint 2

Keep the block input available for the residual addition.

Hint 3

flatten(1) preserves one policy vector and one value per batch item.

Requirements

Use PyTorch float32 computation and outputs.
Residual blocks follow their supplied order; the list may be empty.
Return raw policy logits without softmax or legal-action masking.
Preserve all input tensors and their device.

Constraints

The input has shape $(B,C_{in},N,N)$ .
The trunk convolution weights have shape $(C,C_{in},3,3)$ , and its channel scale and shift each have shape $(C)$ .
Every residual convolution has weight shape $(C,C,3,3)$ , with channel scale and shift vectors of shape $(C)$ .
The policy convolution has weight shape $(C_p,C,1,1)$ , and its linear layer maps $C_pN^2$ values to $N^2+1$ logits.
The value convolution has weight shape $(C_v,C,1,1)$ , its hidden layer maps $C_vN^2$ values to $H$ , and its output layer maps $H$ values to one scalar.
The residual block list may be empty.
Every dimension is positive. All supplied tensor values are finite floating-point values on the same device; input dtypes may differ.
Residual parameters is a list of dictionaries and head parameters is a dictionary, with tensor values under the named keys.
Epsilon is a positive finite number.

Try Similar Problems

x =

Accepts: array

trunk_weight =

Accepts: array

trunk_scale =

Accepts: array

trunk_shift =

Accepts: array

residual_parameters =

Accepts: array

head_parameters =

Accepts: object

epsilon =

Accepts: number