state punishment, update bash scripts & fix logger bugs

Author: Yikai369

sorrel/examples/state_punishment/analysis/analysis.ipynb

@@ -157,7 +157,7 @@ def create_config(
             }
     else:
         social_harm_config = {
-            "A": 5, # 3
+            "A": 10, # 3
             "B": 0, 
             "C": 0, 
             "D": 0.0, 

sorrel/examples/state_punishment/config.py

@@ -5,6 +5,10 @@
 
 from sorrel.utils.logging import ConsoleLogger, TensorboardLogger
 
+# Lowercased entity class names for A–E (see entities.py); always logged, including 0.
+_STATE_PUNISHMENT_RESOURCE_ENCOUNTER_TYPES = ("a", "b", "c", "d", "e")
+_STATE_PUNISHMENT_RESOURCE_ENCOUNTER_SET = frozenset(_STATE_PUNISHMENT_RESOURCE_ENCOUNTER_TYPES)
+
 
 class StatePunishmentLogger:
     """Enhanced logger that tracks encounters and punishment levels."""
@@ -33,6 +37,8 @@ def record_turn(
         encounter_data = {}
 
         if self.multi_agent_env is not None:
+            from sorrel.examples.state_punishment.agents import SeparateModelStatePunishmentAgent
+
             # Initialize total and mean counters
             total_encounters = {}
             mean_encounters = {}
@@ -44,26 +50,39 @@ def record_turn(
             sigma_weights_advantage = []
             sigma_weights_value = []
 
-            # Track action frequencies
-            total_action_frequencies = {}
-            mean_action_frequencies = {}
+            n_envs = len(self.multi_agent_env.individual_envs)
+            # Union of action keys so every epoch logs zeros for unused actions (and vote_* for separate-model agents)
+            global_action_keys: set[str] = set()
+            for env in self.multi_agent_env.individual_envs:
+                ag = env.agents[0]
+                global_action_keys.update(ag.action_names)
+                if isinstance(ag, SeparateModelStatePunishmentAgent):
+                    global_action_keys.update(
+                        ("vote_no", "vote_increase", "vote_decrease")
+                    )
+            sorted_action_keys = sorted(global_action_keys)
 
             for i, env in enumerate(self.multi_agent_env.individual_envs):
                 agent = env.agents[0]
-                agent_count = len(agent.encounters)
 
-                # Individual agent encounter data
+                # Resource encounters (A–E): always log every step, including zeros
+                for res in _STATE_PUNISHMENT_RESOURCE_ENCOUNTER_TYPES:
+                    count = agent.encounters.get(res, 0)
+                    encounter_data[f"Agent_{i}/{res}_encounters"] = count
+                    total_encounters[res] = total_encounters.get(res, 0) + count
+                    mean_encounters[res] = mean_encounters.get(res, 0) + count
+
+                # Other stepped-on entity types (e.g. sand, wall)
                 for entity_type, count in agent.encounters.items():
+                    if entity_type in _STATE_PUNISHMENT_RESOURCE_ENCOUNTER_SET:
+                        continue
                     encounter_data[f"Agent_{i}/{entity_type}_encounters"] = count
-                    
-                    # Initialize if first time seeing this entity type
                     if entity_type not in total_encounters:
                         total_encounters[entity_type] = 0
                         mean_encounters[entity_type] = 0
-                    
                     total_encounters[entity_type] += count
                     mean_encounters[entity_type] += count
-                
+
                 # Individual agent score
                 encounter_data[f"Agent_{i}/individual_score"] = agent.individual_score
                 total_individual_scores += agent.individual_score
@@ -72,21 +91,13 @@ def record_turn(
                 encounter_data[f"Agent_{i}/social_harm_received"] = agent.social_harm_received_epoch
                 total_social_harm_received += agent.social_harm_received_epoch
 
-                # Track action frequencies for this agent
-                for action_name, frequency in agent.action_frequencies.items():
+                # Action frequencies: always emit every key in global union (zeros when unused)
+                for action_name in sorted_action_keys:
+                    frequency = agent.action_frequencies.get(action_name, 0)
                     encounter_data[f"Agent_{i}/action_freq_{action_name}"] = frequency
-                    
-                    # Initialize if first time seeing this action
-                    if action_name not in total_action_frequencies:
-                        total_action_frequencies[action_name] = 0
-                        mean_action_frequencies[action_name] = 0
-                    
-                    total_action_frequencies[action_name] += frequency
-                    mean_action_frequencies[action_name] += frequency
 
                 # Access sigma_weight and epsilon from PyTorchIQN model
                 # Check if agent uses separate models
-                from sorrel.examples.state_punishment.agents import SeparateModelStatePunishmentAgent
                 if isinstance(agent, SeparateModelStatePunishmentAgent):
                     # Separate model agent: log sigma weights and epsilon from both move and vote models
                     # Move model sigma weights
@@ -147,26 +158,32 @@ def record_turn(
             # Add totals and means to encounter_data
             for entity_type in total_encounters:
                 encounter_data[f"Total/total_{entity_type}_encounters"] = total_encounters[entity_type]
-                encounter_data[f"Mean/mean_{entity_type}_encounters"] = mean_encounters[entity_type] / len(self.multi_agent_env.individual_envs)
+                encounter_data[f"Mean/mean_{entity_type}_encounters"] = mean_encounters[entity_type] / n_envs
 
             # Add total and mean individual scores
             encounter_data["Total/total_individual_score"] = total_individual_scores
-            encounter_data["Mean/mean_individual_score"] = total_individual_scores / len(self.multi_agent_env.individual_envs)
+            encounter_data["Mean/mean_individual_score"] = total_individual_scores / n_envs
 
             # Add total and mean social harm received
             encounter_data["Total/total_social_harm_received"] = total_social_harm_received
-            encounter_data["Mean/mean_social_harm_received"] = total_social_harm_received / len(self.multi_agent_env.individual_envs)
+            encounter_data["Mean/mean_social_harm_received"] = total_social_harm_received / n_envs
 
-            # Add total and mean action frequencies
+            # Add total and mean action frequencies (zeros included via .get above)
             # Note: For standard agents, each agent takes one action per turn, so the sum of mean
             # action frequencies should equal max_turns (typically 100) if the epoch completes.
             # For separate model agents, the total includes both movement actions (one per turn)
             # and vote actions (one per vote epoch), so the sum will be higher than max_turns.
             # For example: if max_turns=100 and vote_window_size=10, expect ~110 actions per agent
             # (100 movement + 10 vote actions). Epochs can end early if world.is_done is True.
-            for action_name in total_action_frequencies:
-                encounter_data[f"Total/total_action_freq_{action_name}"] = total_action_frequencies[action_name]
-                encounter_data[f"Mean/mean_action_freq_{action_name}"] = mean_action_frequencies[action_name] / len(self.multi_agent_env.individual_envs)
+            for action_name in sorted_action_keys:
+                total_af = sum(
+                    self.multi_agent_env.individual_envs[j].agents[0].action_frequencies.get(
+                        action_name, 0
+                    )
+                    for j in range(n_envs)
+                )
+                encounter_data[f"Total/total_action_freq_{action_name}"] = total_af
+                encounter_data[f"Mean/mean_action_freq_{action_name}"] = total_af / n_envs
 
 
             # Add mean sigma weights across all agents

sorrel/examples/state_punishment/logger.py

@@ -2,30 +2,32 @@
 # Study 1: same as study 2 but agents can observe punishment level (--punishment_level_accessible).
 # Create four tmux sessions and run CPC experiments (cpc_00/01 x iqn/ppo).
 # Run this from the workspace root (parent of sorrel/), or it will cd there.
+#
+# Long python CLI strings are run via run_cpc_tmux_study1_job.sh because tmux
+# send-keys truncates very long lines (PPO/IQN would stop before --seed 1).
 
 set -e
 ROOT="$(cd "$(dirname "$0")/../../.." && pwd)"
 cd "$ROOT"
 
-# Prefix: ensure conda is available, deactivate any env, then activate sorrel and run
-CONDA_PREFIX='source "$(conda info --base 2>/dev/null)/etc/profile.d/conda.sh" 2>/dev/null; conda deactivate 2>/dev/null; conda deactivate 2>/dev/null; conda activate sorrel && cd '"$ROOT"' && '
+JOB="$ROOT/sorrel/examples/state_punishment/run_cpc_tmux_study1_job.sh"
 
 # cpc_00_ppo: PPO, cpc_weight 0.0
-tmux new-session -d -s study_1_cpc_00_ppo
-tmux send-keys -t study_1_cpc_00_ppo "$CONDA_PREFIX python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --model_type ppo_lstm_cpc --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix np_00_cpc_ppo_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --use_cpc --cpc_horizon 5 --cpc_weight 0.0 --cpc_sample_size 4 --ppo_use_factored_actions --ppo_action_dims \"5,3\" --composite_view --seed 1" C-m
+tmux new-session -d -s study_1_cpc_00_ppo 2>/dev/null || true
+tmux send-keys -t study_1_cpc_00_ppo "bash '$JOB' ppo_00" C-m
 
 # cpc_01_ppo: PPO, cpc_weight 0.1
-tmux new-session -d -s study_1_cpc_01_ppo
-tmux send-keys -t study_1_cpc_01_ppo "$CONDA_PREFIX python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --model_type ppo_lstm_cpc --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix np_01_cpc_ppo_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --use_cpc --cpc_horizon 5 --cpc_weight 0.1 --cpc_sample_size 4 --ppo_use_factored_actions --ppo_action_dims \"5,3\" --composite_view --seed 1" C-m
+tmux new-session -d -s study_1_cpc_01_ppo 2>/dev/null || true
+tmux send-keys -t study_1_cpc_01_ppo "bash '$JOB' ppo_01" C-m
 
 # cpc_01_iqn: IQN, cpc_weight 0.1
-tmux new-session -d -s study_1_cpc_01_iqn
-tmux send-keys -t study_1_cpc_01_iqn "$CONDA_PREFIX python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix np_01_cpc_iqn_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --iqn_use_cpc --iqn_cpc_horizon 5 --iqn_cpc_weight 0.1 --iqn_cpc_sample_size 4 --iqn_use_factored_actions --iqn_action_dims \"5,3\" --composite_views --seed 1" C-m
+tmux new-session -d -s study_1_cpc_01_iqn 2>/dev/null || true
+tmux send-keys -t study_1_cpc_01_iqn "bash '$JOB' iqn_01" C-m
 
 # cpc_00_iqn: IQN, cpc_weight 0.0
-tmux new-session -d -s study_1_cpc_00_iqn
-tmux send-keys -t study_1_cpc_00_iqn "$CONDA_PREFIX python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix np_00_cpc_iqn_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --iqn_use_cpc --iqn_cpc_horizon 5 --iqn_cpc_weight 0.0 --iqn_cpc_sample_size 4 --iqn_use_factored_actions --iqn_action_dims \"5,3\" --composite_views --seed 1" C-m
+tmux new-session -d -s study_1_cpc_00_iqn 2>/dev/null || true
+tmux send-keys -t study_1_cpc_00_iqn "bash '$JOB' iqn_00" C-m
 
-echo "Created 4 tmux sessions: study_1_cpc_00_ppo, study_1_cpc_01_ppo, study_1_cpc_01_iqn, study_1_cpc_00_iqn"
+echo "Ensured 4 tmux sessions exist and sent run commands: study_1_cpc_00_ppo, study_1_cpc_01_ppo, study_1_cpc_01_iqn, study_1_cpc_00_iqn (new-session duplicates ignored)"
 echo "Attach with: tmux attach -t study_1_cpc_00_ppo  (or study_1_cpc_01_ppo, study_1_cpc_01_iqn, study_1_cpc_00_iqn)"
 echo "List sessions: tmux ls"

sorrel/examples/state_punishment/run_cpc_tmux_study1.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# One Study 1 CPC job; invoked by run_cpc_tmux_study1.sh to avoid tmux send-keys length limits.
+set -e
+ROOT="$(cd "$(dirname "$0")/../../.." && pwd)"
+cd "$ROOT"
+# shellcheck source=/dev/null
+source "$(conda info --base 2>/dev/null)/etc/profile.d/conda.sh" 2>/dev/null
+conda deactivate 2>/dev/null || true
+conda deactivate 2>/dev/null || true
+conda activate sorrel
+export PYTHONPATH="$ROOT"
+
+case "$1" in
+  ppo_00)
+    exec python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --model_type ppo_lstm_cpc --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix 00_cpc_ppo_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --use_cpc --cpc_horizon 5 --cpc_weight 0.0 --cpc_sample_size 4 --ppo_use_factored_actions --ppo_action_dims 5,3 --composite_views --seed 1
+    ;;
+  ppo_01)
+    exec python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --model_type ppo_lstm_cpc --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix 01_cpc_ppo_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --use_cpc --cpc_horizon 5 --cpc_weight 0.1 --cpc_sample_size 4 --ppo_use_factored_actions --ppo_action_dims 5,3 --composite_views --seed 1
+    ;;
+  iqn_01)
+    exec python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix 01_cpc_iqn_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --iqn_use_cpc --iqn_cpc_horizon 5 --iqn_cpc_weight 0.1 --iqn_cpc_sample_size 4 --iqn_use_factored_actions --iqn_action_dims 5,3 --composite_views --seed 1
+    ;;
+  iqn_00)
+    exec python sorrel/examples/state_punishment/main.py --num_agents 10 --multi_env_composite --epochs 1500000 --social_harm_accessible --punishment_level_accessible --disable_probe_test --epsilon 0 --use_probabilistic_punishment --use_predefined_punishment_schedule --replacement_selection_mode random_with_tenure --replacement_minimum_tenure_epochs 10000 --agents_to_replace_per_epoch 1 --replacement_start_epoch 10000 --enable_agent_replacement --replacement_min_epochs_between 10000 --run_folder_prefix 00_cpc_iqn_lstm_cpc_slow --respawn_prob 0.01 --num_resources 14 --max_resources 14 --iqn_use_cpc --iqn_cpc_horizon 5 --iqn_cpc_weight 0.0 --iqn_cpc_sample_size 4 --iqn_use_factored_actions --iqn_action_dims 5,3 --composite_views --seed 1
+    ;;
+  *)
+    echo "Usage: $0 {ppo_00|ppo_01|iqn_00|iqn_01}" >&2
+    exit 1
+    ;;
+esac

sorrel/examples/state_punishment/run_cpc_tmux_study1_job.sh

@@ -22,13 +22,13 @@
 
 # increase the punishment prob for A
 # predefined_punishment_probs = np.array([
-#     [0.40, 0.00, 0.00, 0.00, 0.00],  # s = 0
-#     [0.45, 0.05, 0.00, 0.00, 0.00],  # s = 1
+#     [0.30, 0.00, 0.00, 0.00, 0.00],  # s = 0
+#     [0.40, 0.05, 0.00, 0.00, 0.00],  # s = 1
 #     [0.55, 0.10, 0.00, 0.00, 0.00],  # s = 2
 #     [0.70, 0.10, 0.05, 0.00, 0.00],  # s = 3
-#     [0.80, 0.10, 0.10, 0.00, 0.00],  # s = 4
-#     [0.90, 0.10, 0.10, 0.05, 0.00],  # s = 5
-#     [0.95, 0.10, 0.10, 0.10, 0.00],  # s = 6
+#     [0.85, 0.10, 0.10, 0.00, 0.00],  # s = 4
+#     [0.95, 0.10, 0.10, 0.05, 0.00],  # s = 5
+#     [1.0, 0.10, 0.10, 0.10, 0.00],  # s = 6
 #     [1.0, 0.15, 0.10, 0.10, 0.05],  # s = 7
 #     [1.0, 0.15, 0.15, 0.10, 0.10],  # s = 8
 #     [1.0, 0.20, 0.15, 0.15, 0.10],  # s = 9