Custom agents and rewards

An agent selects actions; a reward evaluates each transition during learning. The end-of-episode KPIs evaluate the outcome independently of that reward.

This complete example defines both. Save it as custom_controller.py and run python custom_controller.py. It uses a public dataset and the flat interface.

"""A complete flat-interface example with a custom agent and reward."""

import numpy as np

from citylearn.agents.base import Agent
from citylearn.citylearn import CityLearnEnv
from citylearn.reward_function import RewardFunction


class ZeroActionAgent(Agent):
    """Minimal controller template; replace the action rule for your study."""

    def predict(self, observations, deterministic=None):
        actions = [
            np.clip(np.zeros(space.shape), space.low, space.high).tolist()
            for space in self.action_space
        ]
        self.actions = actions
        self.next_time_step()
        return actions


class GridImportReward(RewardFunction):
    required_observation_names = ("net_electricity_consumption",)

    def calculate(self, observations):
        rewards = [-max(float(o["net_electricity_consumption"]), 0.0) for o in observations]
        return [sum(rewards)] if self.central_agent else rewards


def run(central_agent=False):
    env = CityLearnEnv(
        "citylearn_challenge_2022_phase_all_plus_evs",
        central_agent=central_agent,
        episode_time_steps=24,
        interface="flat",
        reward_function=GridImportReward,
        render_mode="none",
        random_seed=0,
    )
    try:
        agent = ZeroActionAgent(env)
        observations, info = env.reset(seed=0)
        terminated = truncated = False
        while not (terminated or truncated):
            actions = agent.predict(observations)
            observations, rewards, terminated, truncated, info = env.step(actions)
        kpis = env.evaluate_v2(include_business_as_usual=False)
        print(kpis.loc[kpis["cost_function"] == "district_cost_total_control_eur"])
        return kpis
    finally:
        env.close()


if __name__ == "__main__":
    run()

Adapt the action rule

ZeroActionAgent supplies one action vector per action space and preserves the base agent’s action history and timestep. Replace its rule with your controller, using env.action_names and env.observation_names to locate signals. Read bounds from the spaces, rather than assuming all actions use the same range.

Zero actions make this a small coding example, not the operational BAU policy. Use Baselines for that reference. For identified tables and changing asset populations, use Flat and Entity Interfaces.

Adapt the reward

GridImportReward penalizes positive grid imports. It returns one value per building in decentralized mode, or their sum in centralized mode. The required_observation_names attribute lets CityLearn supply only the observations this reward needs.

Pass a reward class directly to CityLearnEnv as shown above, or put an importable class path under reward_function.type in the schema. The module must be importable in the Python environment running the simulation.

Next, inspect Agents, Reward Function and KPIs Reference.