ameya98

Mononofu

# Copyright 2022 DeepMind Technologies Limited.
# Licensed under the Apache License, Version 2.0 and CC BY 4.0.
# You may not use this file except in compliance with these licenses.
# Copies of the licenses can be found at https://www.apache.org/licenses/LICENSE-2.0
# and https://creativecommons.org/licenses/by/4.0/legalcode.

"""Pseudocode description of the Stochastic MuZero algorithm.

This pseudocode was adapted from the original MuZero pseudocode.
"""

ngoodger

import jax.numpy as jnp
from jax import random

class JaxCartPole:
    """
    Based on OpenAI Gym Cartpole
    https://github.com/openai/gym/blob/master/gym/envs/classic_control/cartpole.py
    """
    def __init__(self):
        self.gravity = 9.8

redknightlois

def init_weights(m, variance=1.0):

    def _calculate_fan_in_and_fan_out(tensor):
        dimensions = tensor.dim()
        if dimensions < 2:
            return 1, 1

        if dimensions == 2:  # Linear
            fan_in = tensor.size(1)
            fan_out = tensor.size(0)