Add support for "dodgeball" ; unify some RNG functionality

Author: kcorder

procgen/env.py

@@ -53,8 +53,13 @@
 #################################################
 
 ENV_NAMES_OOD = [
-    "coinrun", 
-    "climber", 
+    "bigfish",
+    "bossfight",
+    "caveflyer",
+    # "chaser",
+    "climber",
+    "coinrun",
+    "dodgeball",
     # Add more as they are finished
 ]
 

procgen/src/basic-abstract-game.cpp

@@ -1121,6 +1121,16 @@ void BasicAbstractGame::choose_random_theme_type_match(const std::shared_ptr<Ent
     }
 }
 
+int BasicAbstractGame::get_random_theme_type_match(const std::shared_ptr<Entity> &ent, const std::string &var_type) {
+    // Get a random theme integer based on the type of variable and train/eval environment. 
+    // Useful for setting multiple entities to the same theme. 
+    if (type_match(var_type)) {
+        return randn_type_switch(asset_num_themes[ent->image_type], var_type);
+    } else {
+        return rand_gen.randn(asset_num_themes[ent->image_type]);
+    }
+}
+
 void BasicAbstractGame::choose_random_theme_switch(const std::shared_ptr<Entity> &ent) {
     // Dispatch method for choosing a random theme 
     if (eval_env) {
@@ -1142,8 +1152,8 @@ void BasicAbstractGame::choose_random_theme_train(const std::shared_ptr<Entity>
         return;
     }
 
-    float ithf = 1 - train_holdout_frac;
-    float theme_frac = num_themes * ithf; 
+    float inv_frac = 1 - train_holdout_frac;
+    float theme_frac = num_themes * inv_frac; 
     int num_train_themes = std::max((int)theme_frac, 1);  // Ensure at least one theme is chosen
     ent->image_theme = rand_gen.randn(num_train_themes);
 
@@ -1170,8 +1180,7 @@ void BasicAbstractGame::choose_random_theme_eval(const std::shared_ptr<Entity> &
         }
         return;
     }
-    float ehf = eval_holdout_frac;
-    float theme_frac = num_themes * ehf; 
+    float theme_frac = num_themes * eval_holdout_frac; 
     int num_eval_themes = std::max((int)theme_frac, 1);
     int start_idx = num_themes - num_eval_themes;
     ent->image_theme = start_idx + rand_gen.randn(num_eval_themes);

procgen/src/basic-abstract-game.h

@@ -112,8 +112,11 @@ class BasicAbstractGame : public Game {
         bool agent_has_collision();
         void reposition_agent();
 
+        void initialize_asset_if_necessary(int img_idx);
+
         // Added for generalization testing -------------------------------------------// 
         bool type_match(std::string var_type); 
+        int get_random_theme_type_match(const std::shared_ptr<Entity> &ent, const std::string &var_type); 
         void choose_random_theme_switch(const std::shared_ptr<Entity> &ent);
         void choose_random_theme_train(const std::shared_ptr<Entity> &ent); 
         void choose_random_theme_eval(const std::shared_ptr<Entity> &ent); 
@@ -180,7 +183,6 @@ class BasicAbstractGame : public Game {
         Grid<int> grid;
 
         QImage *lookup_asset(int img_idx, bool is_reflected = false);
-        void initialize_asset_if_necessary(int img_idx);
         void prepare_for_drawing(float rect_height);
         void draw_background(QPainter &p, const QRect &rect);
         void draw_entity(QPainter &p, const std::shared_ptr<Entity> &to_draw);

procgen/src/games/dodgeball.cpp

@@ -18,7 +18,6 @@ const int DUST_CLOUD = 8;
 const int OOB_WALL = 10;
 
 const int ENEMY_REWARD = 2.0f;
-const int NUM_ENEMY_THEMES = 7;
 
 const float ENEMY_VEL = 0.05f;
 const float BALL_V_ROT = PI * 0.23f;
@@ -348,7 +347,15 @@ class DodgeballGame : public BasicAbstractGame {
 
         spawn_entities(num_enemies, enemy_r, ENEMY, 0, 0, main_width, main_height);
 
-        int enemy_theme = rand_gen.randn(NUM_ENEMY_THEMES);
+        // Choose a random theme for all enemies 
+        int enemy_theme; 
+        for (auto ent : entities) { 
+            if (ent->type == ENEMY) {
+                initialize_asset_if_necessary(ent->image_type);
+                enemy_theme = get_random_theme_type_match(ent, "enemy");
+                break; 
+            } 
+        }
 
         for (auto ent : entities) {
             if (ent->type == ENEMY) {