| 364 | |
| 365 | |
| 366 | class EfficientZeroNet(BaseNet): |
| 367 | def __init__( |
| 368 | self, |
| 369 | observation_shape, |
| 370 | action_space_size, |
| 371 | num_blocks, |
| 372 | num_channels, |
| 373 | reduced_channels_reward, |
| 374 | reduced_channels_value, |
| 375 | reduced_channels_policy, |
| 376 | fc_reward_layers, |
| 377 | fc_value_layers, |
| 378 | fc_policy_layers, |
| 379 | reward_support_size, |
| 380 | value_support_size, |
| 381 | downsample, |
| 382 | inverse_value_transform, |
| 383 | inverse_reward_transform, |
| 384 | lstm_hidden_size, |
| 385 | bn_mt=0.1, |
| 386 | proj_hid=256, |
| 387 | proj_out=256, |
| 388 | pred_hid=64, |
| 389 | pred_out=256, |
| 390 | init_zero=False, |
| 391 | state_norm=False |
| 392 | ): |
| 393 | """EfficientZero network |
| 394 | Parameters |
| 395 | ---------- |
| 396 | observation_shape: tuple or list |
| 397 | shape of observations: [C, W, H] |
| 398 | action_space_size: int |
| 399 | action space |
| 400 | num_blocks: int |
| 401 | number of res blocks |
| 402 | num_channels: int |
| 403 | channels of hidden states |
| 404 | reduced_channels_reward: int |
| 405 | channels of reward head |
| 406 | reduced_channels_value: int |
| 407 | channels of value head |
| 408 | reduced_channels_policy: int |
| 409 | channels of policy head |
| 410 | fc_reward_layers: list |
| 411 | hidden layers of the reward prediction head (MLP head) |
| 412 | fc_value_layers: list |
| 413 | hidden layers of the value prediction head (MLP head) |
| 414 | fc_policy_layers: list |
| 415 | hidden layers of the policy prediction head (MLP head) |
| 416 | reward_support_size: int |
| 417 | dim of reward output |
| 418 | value_support_size: int |
| 419 | dim of value output |
| 420 | downsample: bool |
| 421 | True -> do downsampling for observations. (For board games, do not need) |
| 422 | inverse_value_transform: Any |
| 423 | A function that maps value supports into value scalars |