make rewards configurable

capnspacehook · capnspacehook · commit 84b7d8963779 · 2025-08-20T21:00:43.000-04:00
diff --git a/pufferlib/config/ocean/impulse_wars.ini b/pufferlib/config/ocean/impulse_wars.ini
@@ -30,7 +30,7 @@ continuous = False
 is_training = True
 
 [train]
-total_timesteps = 100_000_000
+total_timesteps = 1_000_000_000
 checkpoint_interval = 250
 
 learning_rate = 0.005
@@ -47,6 +47,78 @@ max = 512
 mean = 128
 scale = auto
 
+# reward parameters
+[sweep.env.reward_win]
+distribution = uniform
+min = 0.0
+mean = 2.0
+max = 5.0
+scale = auto
+
+[sweep.env.reward_self_kill]
+distribution = uniform
+min = -3.0
+mean = -1.0
+max = 0.0
+scale = auto
+
+[sweep.env.reward_enemy_death]
+distribution = uniform
+min = 0.0
+mean = 1.0
+max = 3.0
+scale = auto
+
+[sweep.env.reward_kill]
+distribution = uniform
+min = 0.0
+mean = 1.0
+max = 3.0
+scale = auto
+
+[sweep.env.reward_death]
+distribution = uniform
+min = -1.0
+mean = -0.25
+max = 0.0
+scale = auto
+
+[sweep.env.reward_energy_emptied]
+distribution = uniform
+min = -2.0
+mean = -0.75
+max = 0.0
+scale = auto
+
+[sweep.env.reward_weapon_pickup]
+distribution = uniform
+min = 0.0
+mean = 0.5
+max = 3.0
+scale = auto
+
+[sweep.env.reward_shield_break]
+distribution = uniform
+min = 0.0
+mean = 0.5
+max = 3.0
+scale = auto
+
+[sweep.env.reward_shot_hit_coef]
+distribution = log_normal
+min = 0.0005
+mean = 0.005
+max = 0.05
+scale = auto
+
+[sweep.env.reward_explosion_hit_coef]
+distribution = log_normal
+min = 0.0005
+mean = 0.005
+max = 0.05
+scale = auto
+
+# hyperparameters
 [sweep.train.total_timesteps]
 distribution = log_normal
 min = 250_000_000
diff --git a/pufferlib/ocean/impulse_wars/binding.c b/pufferlib/ocean/impulse_wars/binding.c
@@ -100,6 +100,21 @@ static int my_init(iwEnv *e, PyObject *args, PyObject *kwargs) {
         (bool)unpack(kwargs, "is_training"),
         (bool)unpack(kwargs, "continuous")
     );
+    setRewards(
+        e,
+        (float)unpack(kwargs, "reward_win"),
+        (float)unpack(kwargs, "reward_self_kill"),
+        (float)unpack(kwargs, "reward_enemy_death"),
+        (float)unpack(kwargs, "reward_enemy_kill"),
+        0.0f, // teammate death punishment
+        0.0f, // teammate kill punishment
+        (float)unpack(kwargs, "reward_death"),
+        (float)unpack(kwargs, "reward_energy_emptied"),
+        (float)unpack(kwargs, "reward_weapon_pickup"),
+        (float)unpack(kwargs, "reward_shield_break"),
+        (float)unpack(kwargs, "reward_shot_hit_coef"),
+        (float)unpack(kwargs, "reward_explosion_hit_coef")
+    );
     return 0;
 }
 
@@ -131,6 +146,8 @@ static int my_log(PyObject *dict, Log *log) {
         assign_to_dict(dict, droneLog(buf, i, "total_bursts"), log->stats[i].totalBursts);
         assign_to_dict(dict, droneLog(buf, i, "bursts_hit"), log->stats[i].burstsHit);
         assign_to_dict(dict, droneLog(buf, i, "energy_emptied"), log->stats[i].energyEmptied);
+        assign_to_dict(dict, droneLog(buf, i, "shields_broken"), log->stats[i].shieldsBroken);
+        assign_to_dict(dict, droneLog(buf, i, "own_shield_broken"), log->stats[i].ownShieldBroken);
         assign_to_dict(dict, droneLog(buf, i, "self_kills"), log->stats[i].selfKills);
         assign_to_dict(dict, droneLog(buf, i, "kills"), log->stats[i].kills);
         assign_to_dict(dict, droneLog(buf, i, "wins"), log->stats[i].wins);
diff --git a/pufferlib/ocean/impulse_wars/env.h b/pufferlib/ocean/impulse_wars/env.h
@@ -539,6 +539,19 @@ iwEnv *initEnv(iwEnv *e, uint8_t numDrones, uint8_t numAgents, int8_t mapIdx, ui
     e->sittingDuck = sittingDuck;
     e->isTraining = isTraining;
 
+    e->winReward = WIN_REWARD;
+    e->selfKillPunishment = SELF_KILL_PUNISHMENT;
+    e->enemyDeathReward = ENEMY_DEATH_REWARD;
+    e->enemyKillReward = ENEMY_KILL_REWARD;
+    e->teammateDeathPunishment = TEAMMATE_DEATH_PUNISHMENT;
+    e->teammateKillPunishment = TEAMMATE_KILL_PUNISHMENT;
+    e->deathPunishment = DEATH_PUNISHMENT;
+    e->energyEmptiedPunishment = ENERGY_EMPTY_PUNISHMENT;
+    e->weaponPickupReward = WEAPON_PICKUP_REWARD;
+    e->shieldBreakReward = SHIELD_BREAK_REWARD;
+    e->shotHitRewardCoef = SHOT_HIT_REWARD_COEF;
+    e->explosionHitRewardCoef = EXPLOSION_HIT_REWARD_COEF;
+
     e->obsBytes = obsBytes(e->numDrones);
     e->discreteObsBytes = alignedSize(discreteObsSize(e->numDrones) * sizeof(uint8_t), sizeof(float));
 
@@ -586,6 +599,21 @@ iwEnv *initEnv(iwEnv *e, uint8_t numDrones, uint8_t numAgents, int8_t mapIdx, ui
     return e;
 }
 
+void setRewards(iwEnv *e, float winReward, float selfKillPunishment, float enemyDeathReward, float enemyKillReward, float teammateDeathPunishment, float teammateKillPunishment, float deathPunishment, float energyEmptiedPunishment, float weaponPickupReward, float shieldBreakReward, float shotHitRewardCoef, float explosionHitRewardCoef) {
+    e->winReward = winReward;
+    e->selfKillPunishment = selfKillPunishment;
+    e->enemyDeathReward = enemyDeathReward;
+    e->enemyKillReward = enemyKillReward;
+    e->teammateDeathPunishment = teammateDeathPunishment;
+    e->teammateKillPunishment = teammateKillPunishment;
+    e->deathPunishment = deathPunishment;
+    e->energyEmptiedPunishment = energyEmptiedPunishment;
+    e->weaponPickupReward = weaponPickupReward;
+    e->shieldBreakReward = shieldBreakReward;
+    e->shotHitRewardCoef = shotHitRewardCoef;
+    e->explosionHitRewardCoef = explosionHitRewardCoef;
+}
+
 void clearEnv(iwEnv *e) {
     // rewards get cleared in stepEnv every step
     // memset(e->masks, 1, e->numAgents * sizeof(uint8_t));
@@ -684,15 +712,15 @@ float computeReward(iwEnv *e, droneEntity *drone) {
     float reward = 0.0f;
 
     if (drone->energyFullyDepleted && drone->energyRefillWait == DRONE_ENERGY_REFILL_EMPTY_WAIT) {
-        reward += ENERGY_EMPTY_PUNISHMENT;
+        reward += e->energyEmptiedPunishment;
     }
 
     // only reward picking up a weapon if the standard weapon was
     // previously held; every weapon is better than the standard
     // weapon, but other weapons are situational better so don't
     // reward switching a non-standard weapon
     if (drone->stepInfo.pickedUpWeapon && drone->stepInfo.prevWeapon == STANDARD_WEAPON) {
-        reward += WEAPON_PICKUP_REWARD;
+        reward += e->weaponPickupReward;
     }
 
     for (uint8_t i = 0; i < e->numDrones; i++) {
@@ -704,34 +732,34 @@ float computeReward(iwEnv *e, droneEntity *drone) {
 
         // TODO: punish for hitting teammates?
         if (drone->stepInfo.shotHit[i] != 0.0f && !onTeam) {
-            reward += drone->stepInfo.shotHit[i] * SHOT_HIT_REWARD_COEF;
+            reward += drone->stepInfo.shotHit[i] * e->shotHitRewardCoef;
         }
         if (drone->stepInfo.explosionHit[i] != 0.0f && !onTeam) {
-            reward += drone->stepInfo.explosionHit[i] * EXPLOSION_HIT_REWARD_COEF;
+            reward += drone->stepInfo.explosionHit[i] * e->explosionHitRewardCoef;
         }
         if (drone->stepInfo.brokeShield[i] && !onTeam) {
-            reward += SHIELD_BREAK_REWARD;
+            reward += e->shieldBreakReward;
         }
 
         if (e->numAgents == e->numDrones) {
             if (drone->stepInfo.shotTaken[i] != 0) {
-                reward -= drone->stepInfo.shotTaken[i] * SHOT_HIT_REWARD_COEF;
+                reward -= drone->stepInfo.shotTaken[i] * e->shotHitRewardCoef;
             }
             if (drone->stepInfo.explosionTaken[i]) {
-                reward -= drone->stepInfo.explosionTaken[i] * EXPLOSION_HIT_REWARD_COEF;
+                reward -= drone->stepInfo.explosionTaken[i] * e->explosionHitRewardCoef;
             }
         }
 
         if (enemyDrone->dead && enemyDrone->diedThisStep) {
             if (!onTeam) {
-                reward += ENEMY_DEATH_REWARD;
+                reward += e->enemyDeathReward;
                 if (drone->killed[i]) {
-                    reward += ENEMY_KILL_REWARD;
+                    reward += e->enemyKillReward;
                 }
             } else {
-                reward += TEAMMATE_DEATH_PUNISHMENT;
+                reward += e->teammateDeathPunishment;
                 if (drone->killed[i]) {
-                    reward += TEAMMATE_KILL_PUNISHMENT;
+                    reward += e->teammateKillPunishment;
                 }
             }
             continue;
@@ -756,19 +784,19 @@ const float REWARD_EPS = 1.0e-6f;
 
 void computeRewards(iwEnv *e, const bool roundOver, const int8_t winner, const int8_t winningTeam) {
     if (roundOver && winner != -1 && winner < e->numAgents) {
-        e->rewards[winner] += WIN_REWARD;
+        e->rewards[winner] += e->winReward;
     }
 
     for (uint8_t i = 0; i < e->numDrones; i++) {
         float reward = 0.0f;
         droneEntity *drone = safe_array_get_at(e->drones, i);
         reward = computeReward(e, drone);
         if (!drone->dead && roundOver && winningTeam == drone->team) {
-            reward += WIN_REWARD;
+            reward += e->winReward;
         } else if (drone->diedThisStep) {
-            reward = DEATH_PUNISHMENT;
+            reward = e->deathPunishment;
             if (drone->killedBy == drone->idx) {
-                reward += SELF_KILL_PUNISHMENT;
+                reward += e->selfKillPunishment;
             }
         }
         if (i < e->numAgents) {
@@ -973,6 +1001,10 @@ void addLog(iwEnv *e, Log *log) {
         e->log.stats[j].totalBursts += log->stats[j].totalBursts;
         e->log.stats[j].burstsHit += log->stats[j].burstsHit;
         e->log.stats[j].energyEmptied += log->stats[j].energyEmptied;
+        e->log.stats[j].shieldsBroken += log->stats[j].shieldsBroken;
+        e->log.stats[j].ownShieldBroken += log->stats[j].ownShieldBroken;
+        e->log.stats[j].selfKills += log->stats[j].selfKills;
+        e->log.stats[j].kills += log->stats[j].kills;
 
         for (uint8_t k = 0; k < NUM_WEAPONS; k++) {
             e->log.stats[j].shotsFired[k] += log->stats[j].shotsFired[k];
diff --git a/pufferlib/ocean/impulse_wars/impulse_wars.py b/pufferlib/ocean/impulse_wars/impulse_wars.py
@@ -28,6 +28,16 @@ def __init__(
         continuous: bool = False,
         is_training: bool = True,
         human_control: bool = False,
+        reward_win: float = 2.0,
+        reward_self_kill: float = -1.0,
+        reward_enemy_death: float = 1.0,
+        reward_enemy_kill: float = 1.0,
+        reward_death: float = -0.25,
+        reward_energy_emptied: float = -0.75,
+        reward_weapon_pickup: float = 0.5,
+        reward_shield_break: float = 0.5,
+        reward_shot_hit_coef: float = 0.005, 
+        reward_explosion_hit_coef: float = 0.005,
         seed: int = 0,
         render: bool = False,
         report_interval: int = 64,
@@ -98,6 +108,16 @@ def __init__(
             sitting_duck=sitting_duck,
             is_training=is_training,
             continuous=continuous,
+            reward_win=reward_win,
+            reward_self_kill=reward_self_kill,
+            reward_enemy_death=reward_enemy_death,
+            reward_enemy_kill=reward_enemy_kill,
+            reward_death=reward_death,
+            reward_energy_emptied=reward_energy_emptied,
+            reward_weapon_pickup=reward_weapon_pickup,
+            reward_shield_break=reward_shield_break,
+            reward_shot_hit_coef=reward_shot_hit_coef,
+            reward_explosion_hit_coef=reward_explosion_hit_coef,
         )
 
         binding.shared(self.c_envs)
diff --git a/pufferlib/ocean/impulse_wars/types.h b/pufferlib/ocean/impulse_wars/types.h
@@ -412,6 +412,19 @@ typedef struct iwEnv {
     bool sittingDuck;
     bool isTraining;
 
+    float winReward;
+    float selfKillPunishment;
+    float enemyDeathReward;
+    float enemyKillReward;
+    float teammateDeathPunishment;
+    float teammateKillPunishment;
+    float deathPunishment;
+    float energyEmptiedPunishment;
+    float weaponPickupReward;
+    float shieldBreakReward;
+    float shotHitRewardCoef;
+    float explosionHitRewardCoef;
+
     uint16_t obsBytes;
     uint16_t discreteObsBytes;
     bool continuousActions;
