MuJoCo gym partial port

.gitignore

@@ -178,6 +178,7 @@ resources/drive/data/*
 resources/drive/binaries/*
 resources/boxoban/levels/
 resources/boxoban/boxoban_maps_*.bin
+resources/mujoco/*.bin
 
 # Policy weights live in the website repo (docs/assets/models/)
 resources/**/*_weights.bin

build.sh

@@ -13,6 +13,8 @@ set -e
 #   ./build.sh breakout --web        # Emscripten web build
 #                                    # copy build/web/ENV/* to ../docker/puffer.ai/docs/assets/ENV/
 #   ./build.sh breakout --profile    # Kernel profiling binary
+#   ./build.sh mjc_half_cheetah      # MuJoCo-family envs live in ocean/mujoco/ENV.h
+#   ./build.sh mjc_half_cheetah --cu # ... and ocean/mujoco/ENV.cu (one GPU thread per env)
 #   ./build.sh all                   # Build all envs native and native float32
 #
 # Env is compiled in. Run: ./puffer train|eval|match|sweep [--section.key=value ...]
@@ -182,6 +184,8 @@ elif [ "$ENV" = "nethack" ]; then
                     -Xlinker -rpath -Xlinker "$NETHACK_LIB_DIR" -ldl)
 elif [ -d "ocean/$ENV" ]; then
     SRC_DIR="ocean/$ENV"
+elif [ -f "ocean/mujoco/$ENV.h" ]; then
+    SRC_DIR="ocean/mujoco"
 else
     echo "Error: environment '$ENV' not found" && exit 1
 fi

config/mjc_ant.ini

@@ -0,0 +1,107 @@
+[base]
+env_name = mjc_ant
+
+[vec]
+total_agents = 256
+num_buffers = 1
+num_threads = 16
+
+[env]
+model = resources/mujoco/ant.bin
+max_steps = 1000
+reset_noise_scale = 0.1
+forward_reward_weight = 1.0
+ctrl_cost_weight = 0.5
+contact_cost_weight = 0.0005
+healthy_reward = 1.0
+
+[policy]
+hidden_size = 512
+num_layers = 2
+expansion_factor = 1
+
+[train]
+gpus = 1
+total_timesteps = 10000000
+learning_rate = 0.003
+anneal_lr = 1
+min_lr_ratio = 0
+gamma = 0.978
+gae_lambda = 0.944
+replay_ratio = 5.7
+clip_coef = 0.23
+vf_coef = 1.6
+vf_clip_coef = 1.0
+max_grad_norm = 1.05
+ent_coef = 0.00031
+momentum = 0.99
+minibatch_size = 2048
+horizon = 32
+vtrace_rho_clip = 1.0
+vtrace_c_clip = 1.0
+
+[sweep]
+metric = score
+goal = maximize
+max_runs = 30
+
+[sweep.train.total_timesteps]
+min = 1e7
+max = 1.2e7
+
+[sweep.policy.hidden_size]
+min = 256
+max = 512
+
+[sweep.policy.num_layers]
+distribution = int_uniform
+min = 2
+max = 3
+
+[sweep.vec.total_agents]
+min = 128
+max = 1024
+
+[sweep.train.horizon]
+min = 32
+max = 256
+
+[sweep.train.minibatch_size]
+min = 512
+max = 4096
+
+[sweep.train.replay_ratio]
+min = 2
+max = 6
+
+[sweep.train.learning_rate]
+min = 0.0005
+max = 0.02
+
+[sweep.train.momentum]
+min = 0.5
+max = 0.99
+
+[sweep.train.ent_coef]
+min = 0.00001
+max = 0.01
+
+[sweep.train.gamma]
+min = 0.88
+max = 0.99
+
+[sweep.train.gae_lambda]
+min = 0.8
+max = 0.95
+
+[sweep.train.clip_coef]
+min = 0.1
+max = 0.35
+
+[sweep.train.vf_coef]
+min = 0.5
+max = 2.0
+
+[sweep.train.vf_clip_coef]
+min = 1.0
+max = 5.0

config/mjc_half_cheetah.ini

@@ -0,0 +1,105 @@
+[base]
+env_name = mjc_half_cheetah
+
+[vec]
+total_agents = 512
+num_buffers = 1
+num_threads = 16
+
+[env]
+model = resources/mujoco/half_cheetah.bin
+max_steps = 1000
+reset_noise_scale = 0.1
+forward_reward_weight = 1.0
+ctrl_cost_weight = 0.1
+
+[policy]
+hidden_size = 256
+num_layers = 2
+expansion_factor = 1
+
+[train]
+gpus = 1
+total_timesteps = 10000000
+learning_rate = 0.0036
+anneal_lr = 1
+min_lr_ratio = 0
+gamma = 0.975
+gae_lambda = 0.922
+replay_ratio = 4.9
+clip_coef = 0.11
+vf_coef = 0.5
+vf_clip_coef = 4.1
+max_grad_norm = 1.2
+ent_coef = 0.0017
+momentum = 0.955
+minibatch_size = 1024
+horizon = 32
+vtrace_rho_clip = 1.0
+vtrace_c_clip = 1.0
+
+[sweep]
+metric = score
+goal = maximize
+max_runs = 30
+
+[sweep.train.total_timesteps]
+min = 1e7
+max = 1.2e7
+
+[sweep.policy.hidden_size]
+min = 256
+max = 512
+
+[sweep.policy.num_layers]
+distribution = int_uniform
+min = 2
+max = 3
+
+[sweep.vec.total_agents]
+min = 128
+max = 1024
+
+[sweep.train.horizon]
+min = 32
+max = 256
+
+[sweep.train.minibatch_size]
+min = 512
+max = 4096
+
+[sweep.train.replay_ratio]
+min = 2
+max = 6
+
+[sweep.train.learning_rate]
+min = 0.0005
+max = 0.02
+
+[sweep.train.momentum]
+min = 0.5
+max = 0.99
+
+[sweep.train.ent_coef]
+min = 0.00001
+max = 0.01
+
+[sweep.train.gamma]
+min = 0.88
+max = 0.99
+
+[sweep.train.gae_lambda]
+min = 0.8
+max = 0.95
+
+[sweep.train.clip_coef]
+min = 0.1
+max = 0.35
+
+[sweep.train.vf_coef]
+min = 0.5
+max = 2.0
+
+[sweep.train.vf_clip_coef]
+min = 1.0
+max = 5.0

config/mjc_hopper.ini

@@ -0,0 +1,106 @@
+[base]
+env_name = mjc_hopper
+
+[vec]
+total_agents = 256
+num_buffers = 1
+num_threads = 16
+
+[env]
+model = resources/mujoco/hopper.bin
+max_steps = 1000
+reset_noise_scale = 0.005
+forward_reward_weight = 1.0
+ctrl_cost_weight = 0.001
+healthy_reward = 1.0
+
+[policy]
+hidden_size = 256
+num_layers = 2
+expansion_factor = 1
+
+[train]
+gpus = 1
+total_timesteps = 10000000
+learning_rate = 0.0017
+anneal_lr = 1
+min_lr_ratio = 0
+gamma = 0.983
+gae_lambda = 0.939
+replay_ratio = 4.5
+clip_coef = 0.26
+vf_coef = 1.8
+vf_clip_coef = 2.6
+max_grad_norm = 0.97
+ent_coef = 0.00096
+momentum = 0.974
+minibatch_size = 512
+horizon = 128
+vtrace_rho_clip = 1.0
+vtrace_c_clip = 1.0
+
+[sweep]
+metric = score
+goal = maximize
+max_runs = 30
+
+[sweep.train.total_timesteps]
+min = 1e7
+max = 1.2e7
+
+[sweep.policy.hidden_size]
+min = 256
+max = 512
+
+[sweep.policy.num_layers]
+distribution = int_uniform
+min = 2
+max = 3
+
+[sweep.vec.total_agents]
+min = 128
+max = 1024
+
+[sweep.train.horizon]
+min = 32
+max = 256
+
+[sweep.train.minibatch_size]
+min = 512
+max = 4096
+
+[sweep.train.replay_ratio]
+min = 2
+max = 6
+
+[sweep.train.learning_rate]
+min = 0.0005
+max = 0.02
+
+[sweep.train.momentum]
+min = 0.5
+max = 0.99
+
+[sweep.train.ent_coef]
+min = 0.00001
+max = 0.01
+
+[sweep.train.gamma]
+min = 0.88
+max = 0.99
+
+[sweep.train.gae_lambda]
+min = 0.8
+max = 0.95
+
+[sweep.train.clip_coef]
+min = 0.1
+max = 0.35
+
+[sweep.train.vf_coef]
+min = 0.5
+max = 2.0
+
+[sweep.train.vf_clip_coef]
+min = 1.0
+max = 5.0