fix IW training on multiple GPUs, removed unused param

Author: capnspacehook

pufferlib/ocean/impulse_wars/scripted_agent.h

@@ -343,7 +343,7 @@ void handleWallProximity(iwEnv *e, const droneEntity *drone, const wallEntity *w
 
 // charge burst until we're close enough to a death wall to burst off
 // of it
-void wallBurst(iwEnv *e, const droneEntity *drone, const float speed, const float distance, agentActions *actions) {
+void wallBurst(iwEnv *e, const droneEntity *drone, const float distance, agentActions *actions) {
     if (distance < DRONE_BURST_RADIUS_MIN) {
         scriptedAgentBurst(drone, actions);
         return;
@@ -432,7 +432,7 @@ agentActions scriptedAgentActions(iwEnv *e, droneEntity *drone) {
             if (entityTypeIsWall(ent->type) && ent->type == DEATH_WALL_ENTITY) {
                 actions.brake = true;
                 if (drone->shield == NULL) {
-                    wallBurst(e, drone, droneSpeed, b2Distance(drone->pos, ctx.point), &actions);
+                    wallBurst(e, drone, b2Distance(drone->pos, ctx.point), &actions);
                 }
 
                 const b2Vec2 droneDirection = b2Normalize(drone->velocity);

pufferlib/ocean/torch.py

@@ -598,7 +598,6 @@ def __init__(
         num_drones: int = 2,
         continuous: bool = False,
         is_training: bool = True,
-        device: str = "cuda",
         **kwargs,
     ):
         super().__init__()
@@ -616,13 +615,13 @@ def __init__(
             + [self.obsInfo.wallTypes + 1] * self.obsInfo.numFloatingWallObs
             + [self.numDrones + 1] * self.obsInfo.numProjectileObs,
         )
-        discreteOffsets = torch.tensor([0] + list(np.cumsum(self.discreteFactors)[:-1]), device=device).view(
+        discreteOffsets = torch.tensor([0] + list(np.cumsum(self.discreteFactors)[:-1])).view(
             1, -1
         )
         self.register_buffer("discreteOffsets", discreteOffsets, persistent=False)
         self.discreteMultihotDim = self.discreteFactors.sum()
 
-        multihotBuffer = torch.zeros(batch_size, self.discreteMultihotDim, device=device)
+        multihotBuffer = torch.zeros(batch_size, self.discreteMultihotDim)
         self.register_buffer("multihotOutput", multihotBuffer, persistent=False)
 
         # most of the observation is a 2D array of bytes, but the end