Spaces:

Shreya2005-2005
/

study-session-optimizer

Sleeping

study-session-optimizer / src /environment.py

Shreya2005-2005

FINAL FIX: indentation fix and hardened scores to [0.1, 0.9]

cc5d66e 7 months ago

11.9 kB

# environment.py
# The actual RL environment logic for the Smart Study Session Optimizer
#
# The core idea: student has energy + focus, each action moves time forward by
# a fixed interval. Energy drains while studying, recovers on breaks.
# Focus tracks how well you can concentrate.
#
# NOTE: this works, don't touch it (the reward shaping took forever to tune)

import random
import math
from typing import Optional
from .models import Observation, StepResult

# Time constants (in minutes)
STUDY_INTERVAL = 5        # each "study" action = 5 minutes
SHORT_BREAK = 10          # short break duration
LONG_BREAK = 20           # long break duration

# Energy mechanics
ENERGY_DRAIN_BASE = 3.0       # energy lost per study interval
ENERGY_DRAIN_FOCUS_BONUS = 1.5  # extra drain when you're really focused
SHORT_BREAK_RECOVERY = 15.0
LONG_BREAK_RECOVERY = 35.0

# Focus mechanics - focus degrades the longer you study without a break
FOCUS_DECAY_RATE = 2.5        # focus lost per study interval
FOCUS_RECOVERY_SHORT = 20.0
FOCUS_RECOVERY_LONG = 45.0

# Thresholds
LOW_ENERGY_THRESHOLD = 40.0   # below this = bad idea to keep studying
HIGH_FOCUS_THRESHOLD = 60.0   # above this = productive


class StudyEnvironment:
    """
    The main environment class. Simulates a student's study session.
    State space:
        - energy (0-100): physical/mental stamina
        - focus (0-100): concentration level right now
        - time_studied: cumulative study minutes
        - session_minutes: total session time elapsed
        - consecutive_study_minutes: streak without a break
        - break_count: breaks taken
        - productive_minutes: minutes with focus > 60
    Action space:
        0 = keep studying (5 min)
        1 = short break (10 min)
        2 = long break (20 min)
    """

    @staticmethod
    def _strict_score(value: float) -> float:
        """Clamp a score strictly into [0.1, 0.9].
        The validator rejects 0.0 and 1.0, so we use a very safe margin."""
        return max(0.1, min(0.9, round(float(value), 4)))

    def __init__(self, task_id: str = "maximize_productivity", max_session_minutes: int = 480):
        self.task_id = task_id
        self.max_session_minutes = max_session_minutes

        # task-specific overrides
        # TODO: maybe make this a config dict later
        if task_id == "survive_8_hours":
            self.max_session_minutes = 480
        elif task_id == "maximize_productivity":
            self.max_session_minutes = 240
        elif task_id == "smart_break_taker":
            self.max_session_minutes = 400

        self._state = {}
        self.reset()

    def reset(self) -> Observation:
        """Start fresh. Called at the beginning of each episode."""
        print(f"[ENV] Resetting environment, task={self.task_id}")

        # small random variation so the agent has to generalize
        e = 85.0 + random.uniform(-10, 10)   # e = energy (yeah, short name, deal with it)
        f = 80.0 + random.uniform(-10, 10)   # f = focus

        self._state = {
            "energy": min(100.0, max(0.0, e)),
            "focus": min(100.0, max(0.0, f)),
            "time_studied": 0.0,
            "session_minutes": 0.0,
            "consecutive_study_minutes": 0.0,
            "break_count": 0,
            "productive_minutes": 0.0,
            # tracking stuff for graders
            "good_breaks": 0,   # breaks taken when energy < 40
            "total_breaks": 0,
            "done": False,
            "step_count": 0,
        }

        return self._get_obs()

    def step(self, action: int) -> StepResult:
        """
        Execute one action and return the result.
        action: 0=study, 1=short_break, 2=long_break
        """
        if self._state["done"]:
            # this shouldn't happen but just in case
            print("[ENV] WARNING: step() called after episode done")
            return StepResult(
                observation=self._get_obs(),
                reward=0.0,
                done=True,
                info=self._get_info()
            )

        s = self._state   # shorthand (this works, don't touch it)
        s["step_count"] += 1

        reward = 0.0

        if action == 0:  # study
            reward = self._do_study()
        elif action == 1:  # short break
            reward = self._do_short_break()
        elif action == 2:  # long break
            reward = self._do_long_break()
        else:
            print(f"[ENV] Unknown action {action}, defaulting to study")
            reward = self._do_study()

        # check termination conditions
        done = self._check_done()
        s["done"] = done

        if done:
            # add a terminal reward based on task
            reward += self._terminal_reward()

        obs = self._get_obs()
        info = self._get_info()

        print(f"[ENV] Step {s['step_count']}: action={action}, reward={reward:.3f}, "
              f"energy={s['energy']:.1f}, focus={s['focus']:.1f}, done={done}")

        return StepResult(observation=obs, reward=reward, done=done, info=info)

    def _do_study(self) -> float:
        """Handle the study action. Returns reward for this step."""
        s = self._state

        # energy drains faster when you're really focused (makes sense actually)
        focus_factor = s["focus"] / 100.0
        drain = ENERGY_DRAIN_BASE + (ENERGY_DRAIN_FOCUS_BONUS * focus_factor)
        s["energy"] = max(0.0, s["energy"] - drain)

        # focus decays the longer you study without a break
        consecutive_factor = min(s["consecutive_study_minutes"] / 60.0, 1.5)
        focus_drop = FOCUS_DECAY_RATE * (1 + consecutive_factor * 0.5)
        s["focus"] = max(0.0, s["focus"] - focus_drop)

        # advance time
        s["time_studied"] += STUDY_INTERVAL
        s["session_minutes"] += STUDY_INTERVAL
        s["consecutive_study_minutes"] += STUDY_INTERVAL

        # check if this was a productive interval
        if s["focus"] >= HIGH_FOCUS_THRESHOLD:
            s["productive_minutes"] += STUDY_INTERVAL

        # reward calculation
        reward = 0.0

        # positive reward for studying when you have energy and focus
        if s["energy"] > LOW_ENERGY_THRESHOLD and s["focus"] > HIGH_FOCUS_THRESHOLD:
            reward += 1.0
        elif s["energy"] > LOW_ENERGY_THRESHOLD:
            reward += 0.3
        else:
            # studying while exhausted - penalize slightly
            reward -= 0.5

        # small bonus just for getting through time (survival reward)
        reward += 0.1

        return reward

    def _do_short_break(self) -> float:
        """Short break: 10 minutes, moderate recovery."""
        s = self._state

        e_before = s["energy"]
        s["total_breaks"] += 1
        s["break_count"] += 1

        # check if this was a "smart" break (energy was low)
        if e_before < LOW_ENERGY_THRESHOLD:
            s["good_breaks"] += 1

        s["energy"] = min(100.0, s["energy"] + SHORT_BREAK_RECOVERY)
        s["focus"] = min(100.0, s["focus"] + FOCUS_RECOVERY_SHORT)
        s["consecutive_study_minutes"] = 0.0

        # time passes but no studying happens
        s["session_minutes"] += SHORT_BREAK

        # reward: good if you needed the break, meh if you didn't
        if e_before < LOW_ENERGY_THRESHOLD:
            reward = 0.8   # smart move
        elif e_before < 60.0:
            reward = 0.2
        else:
            reward = -0.3  # wasted time when you could have been studying

        return reward

    def _do_long_break(self) -> float:
        """Long break: 20 minutes, full recovery but costs more time."""
        s = self._state

        e_before = s["energy"]
        s["total_breaks"] += 1
        s["break_count"] += 1

        if e_before < LOW_ENERGY_THRESHOLD:
            s["good_breaks"] += 1

        s["energy"] = min(100.0, s["energy"] + LONG_BREAK_RECOVERY)
        s["focus"] = min(100.0, s["focus"] + FOCUS_RECOVERY_LONG)
        s["consecutive_study_minutes"] = 0.0
        s["session_minutes"] += LONG_BREAK

        # long break is only a good call when you're really drained
        if e_before < 25.0:
            reward = 1.0   # definitely needed this
        elif e_before < LOW_ENERGY_THRESHOLD:
            reward = 0.5
        else:
            reward = -0.6  # overkill, wasted too much time

        return reward

    def _check_done(self) -> bool:
        """Is the episode over?"""
        s = self._state

        # ran out of energy - bad
        if s["energy"] <= 0.0:
            print("[ENV] Episode ended: energy depleted!")
            return True

        # reached max session time
        if s["session_minutes"] >= self.max_session_minutes:
            print(f"[ENV] Episode ended: session complete ({s['session_minutes']} min)")
            return True

        # TODO: add a step limit as safety net for infinite loops

        return False

    def _terminal_reward(self) -> float:
        """Bonus/penalty given at the very end of the episode."""
        s = self._state
        terminal = 0.0

        if self.task_id == "survive_8_hours":
            # reward for surviving with energy left
            terminal = (s["energy"] / 100.0) * 5.0
        elif self.task_id == "maximize_productivity":
            # big reward for productive minutes
            terminal = (s["productive_minutes"] / 240.0) * 5.0
        elif self.task_id == "smart_break_taker":
            if s["total_breaks"] > 0:
                ratio = s["good_breaks"] / s["total_breaks"]
                terminal = ratio * 5.0

        return terminal

    def _get_obs(self) -> Observation:
        """Build the current observation from internal state."""
        s = self._state
        return Observation(
            energy=round(s["energy"], 2),
            focus=round(s["focus"], 2),
            time_studied=round(s["time_studied"], 2),
            session_minutes=round(s["session_minutes"], 2),
            consecutive_study_minutes=round(s["consecutive_study_minutes"], 2),
            break_count=s["break_count"],
            productive_minutes=round(s["productive_minutes"], 2),
        )

    def _get_info(self) -> dict:
        """Extra info for graders and debugging.
        CRITICAL: every score MUST be strictly in (0, 1) — never 0.0 or 1.0."""
        s = self._state
        total_breaks = max(s["total_breaks"], 1)

        # raw scores
        survive_score = s["energy"] / 100.0
        productivity_score = s["productive_minutes"] / 240.0

        if total_breaks > 0:
            break_score = s["good_breaks"] / total_breaks
        else:
            break_score = 0.0

        # strictly clamp into (0,1) — _strict_score guarantees (0.01, 0.99)
        scores = {
            "survive_8_hours": self._strict_score(survive_score),
            "maximize_productivity": self._strict_score(productivity_score),
            "smart_break_taker": self._strict_score(break_score),
        }

        current_score = scores.get(self.task_id, 0.5)

        # FINAL SAFETY NET: if anything is still 0.0 or 1.0, fix it
        for k in scores:
            if scores[k] <= 0.05:
                scores[k] = 0.1
            if scores[k] >= 0.95:
                scores[k] = 0.9
        if current_score <= 0.05:
            current_score = 0.1
        if current_score >= 0.95:
            current_score = 0.9

        return {
            "task_id": self.task_id,
            "score": current_score,
            "step_count": s["step_count"],
            "total_breaks": s["total_breaks"],
            "good_breaks": s["good_breaks"],
            "energy_remaining": round(s["energy"], 2),
            "productive_minutes": round(s["productive_minutes"], 2),
            "all_task_scores": scores,
        }

    def get_state(self) -> dict:
        """Return full internal state (for the /state endpoint)."""
        return {**self._state, "task_id": self.task_id}

	# environment.py
	# The actual RL environment logic for the Smart Study Session Optimizer
	#
	# The core idea: student has energy + focus, each action moves time forward by
	# a fixed interval. Energy drains while studying, recovers on breaks.
	# Focus tracks how well you can concentrate.
	#
	# NOTE: this works, don't touch it (the reward shaping took forever to tune)

	import random
	import math
	from typing import Optional
	from .models import Observation, StepResult

	# Time constants (in minutes)
	STUDY_INTERVAL = 5 # each "study" action = 5 minutes
	SHORT_BREAK = 10 # short break duration
	LONG_BREAK = 20 # long break duration

	# Energy mechanics
	ENERGY_DRAIN_BASE = 3.0 # energy lost per study interval
	ENERGY_DRAIN_FOCUS_BONUS = 1.5 # extra drain when you're really focused
	SHORT_BREAK_RECOVERY = 15.0
	LONG_BREAK_RECOVERY = 35.0

	# Focus mechanics - focus degrades the longer you study without a break
	FOCUS_DECAY_RATE = 2.5 # focus lost per study interval
	FOCUS_RECOVERY_SHORT = 20.0
	FOCUS_RECOVERY_LONG = 45.0

	# Thresholds
	LOW_ENERGY_THRESHOLD = 40.0 # below this = bad idea to keep studying
	HIGH_FOCUS_THRESHOLD = 60.0 # above this = productive


	class StudyEnvironment:
	"""
	The main environment class. Simulates a student's study session.

	State space:
	- energy (0-100): physical/mental stamina
	- focus (0-100): concentration level right now
	- time_studied: cumulative study minutes
	- session_minutes: total session time elapsed
	- consecutive_study_minutes: streak without a break
	- break_count: breaks taken
	- productive_minutes: minutes with focus > 60

	Action space:
	0 = keep studying (5 min)
	1 = short break (10 min)
	2 = long break (20 min)
	"""

	@staticmethod
	def _strict_score(value: float) -> float:
	"""Clamp a score strictly into [0.1, 0.9].
	The validator rejects 0.0 and 1.0, so we use a very safe margin."""
	return max(0.1, min(0.9, round(float(value), 4)))

	def __init__(self, task_id: str = "maximize_productivity", max_session_minutes: int = 480):
	self.task_id = task_id
	self.max_session_minutes = max_session_minutes

	# task-specific overrides
	# TODO: maybe make this a config dict later
	if task_id == "survive_8_hours":
	self.max_session_minutes = 480
	elif task_id == "maximize_productivity":
	self.max_session_minutes = 240
	elif task_id == "smart_break_taker":
	self.max_session_minutes = 400

	self._state = {}
	self.reset()

	def reset(self) -> Observation:
	"""Start fresh. Called at the beginning of each episode."""
	print(f"[ENV] Resetting environment, task={self.task_id}")

	# small random variation so the agent has to generalize
	e = 85.0 + random.uniform(-10, 10) # e = energy (yeah, short name, deal with it)
	f = 80.0 + random.uniform(-10, 10) # f = focus

	self._state = {
	"energy": min(100.0, max(0.0, e)),
	"focus": min(100.0, max(0.0, f)),
	"time_studied": 0.0,
	"session_minutes": 0.0,
	"consecutive_study_minutes": 0.0,
	"break_count": 0,
	"productive_minutes": 0.0,
	# tracking stuff for graders
	"good_breaks": 0, # breaks taken when energy < 40
	"total_breaks": 0,
	"done": False,
	"step_count": 0,
	}

	return self._get_obs()

	def step(self, action: int) -> StepResult:
	"""
	Execute one action and return the result.
	action: 0=study, 1=short_break, 2=long_break
	"""
	if self._state["done"]:
	# this shouldn't happen but just in case
	print("[ENV] WARNING: step() called after episode done")
	return StepResult(
	observation=self._get_obs(),
	reward=0.0,
	done=True,
	info=self._get_info()
	)

	s = self._state # shorthand (this works, don't touch it)
	s["step_count"] += 1

	reward = 0.0

	if action == 0: # study
	reward = self._do_study()
	elif action == 1: # short break
	reward = self._do_short_break()
	elif action == 2: # long break
	reward = self._do_long_break()
	else:
	print(f"[ENV] Unknown action {action}, defaulting to study")
	reward = self._do_study()

	# check termination conditions
	done = self._check_done()
	s["done"] = done

	if done:
	# add a terminal reward based on task
	reward += self._terminal_reward()

	obs = self._get_obs()
	info = self._get_info()

	print(f"[ENV] Step {s['step_count']}: action={action}, reward={reward:.3f}, "
	f"energy={s['energy']:.1f}, focus={s['focus']:.1f}, done={done}")

	return StepResult(observation=obs, reward=reward, done=done, info=info)

	def _do_study(self) -> float:
	"""Handle the study action. Returns reward for this step."""
	s = self._state

	# energy drains faster when you're really focused (makes sense actually)
	focus_factor = s["focus"] / 100.0
	drain = ENERGY_DRAIN_BASE + (ENERGY_DRAIN_FOCUS_BONUS * focus_factor)
	s["energy"] = max(0.0, s["energy"] - drain)

	# focus decays the longer you study without a break
	consecutive_factor = min(s["consecutive_study_minutes"] / 60.0, 1.5)
	focus_drop = FOCUS_DECAY_RATE * (1 + consecutive_factor * 0.5)
	s["focus"] = max(0.0, s["focus"] - focus_drop)

	# advance time
	s["time_studied"] += STUDY_INTERVAL
	s["session_minutes"] += STUDY_INTERVAL
	s["consecutive_study_minutes"] += STUDY_INTERVAL

	# check if this was a productive interval
	if s["focus"] >= HIGH_FOCUS_THRESHOLD:
	s["productive_minutes"] += STUDY_INTERVAL

	# reward calculation
	reward = 0.0

	# positive reward for studying when you have energy and focus
	if s["energy"] > LOW_ENERGY_THRESHOLD and s["focus"] > HIGH_FOCUS_THRESHOLD:
	reward += 1.0
	elif s["energy"] > LOW_ENERGY_THRESHOLD:
	reward += 0.3
	else:
	# studying while exhausted - penalize slightly
	reward -= 0.5

	# small bonus just for getting through time (survival reward)
	reward += 0.1

	return reward

	def _do_short_break(self) -> float:
	"""Short break: 10 minutes, moderate recovery."""
	s = self._state

	e_before = s["energy"]
	s["total_breaks"] += 1
	s["break_count"] += 1

	# check if this was a "smart" break (energy was low)
	if e_before < LOW_ENERGY_THRESHOLD:
	s["good_breaks"] += 1

	s["energy"] = min(100.0, s["energy"] + SHORT_BREAK_RECOVERY)
	s["focus"] = min(100.0, s["focus"] + FOCUS_RECOVERY_SHORT)
	s["consecutive_study_minutes"] = 0.0

	# time passes but no studying happens
	s["session_minutes"] += SHORT_BREAK

	# reward: good if you needed the break, meh if you didn't
	if e_before < LOW_ENERGY_THRESHOLD:
	reward = 0.8 # smart move
	elif e_before < 60.0:
	reward = 0.2
	else:
	reward = -0.3 # wasted time when you could have been studying

	return reward

	def _do_long_break(self) -> float:
	"""Long break: 20 minutes, full recovery but costs more time."""
	s = self._state

	e_before = s["energy"]
	s["total_breaks"] += 1
	s["break_count"] += 1

	if e_before < LOW_ENERGY_THRESHOLD:
	s["good_breaks"] += 1

	s["energy"] = min(100.0, s["energy"] + LONG_BREAK_RECOVERY)
	s["focus"] = min(100.0, s["focus"] + FOCUS_RECOVERY_LONG)
	s["consecutive_study_minutes"] = 0.0
	s["session_minutes"] += LONG_BREAK

	# long break is only a good call when you're really drained
	if e_before < 25.0:
	reward = 1.0 # definitely needed this
	elif e_before < LOW_ENERGY_THRESHOLD:
	reward = 0.5
	else:
	reward = -0.6 # overkill, wasted too much time

	return reward

	def _check_done(self) -> bool:
	"""Is the episode over?"""
	s = self._state

	# ran out of energy - bad
	if s["energy"] <= 0.0:
	print("[ENV] Episode ended: energy depleted!")
	return True

	# reached max session time
	if s["session_minutes"] >= self.max_session_minutes:
	print(f"[ENV] Episode ended: session complete ({s['session_minutes']} min)")
	return True

	# TODO: add a step limit as safety net for infinite loops

	return False

	def _terminal_reward(self) -> float:
	"""Bonus/penalty given at the very end of the episode."""
	s = self._state
	terminal = 0.0

	if self.task_id == "survive_8_hours":
	# reward for surviving with energy left
	terminal = (s["energy"] / 100.0) * 5.0
	elif self.task_id == "maximize_productivity":
	# big reward for productive minutes
	terminal = (s["productive_minutes"] / 240.0) * 5.0
	elif self.task_id == "smart_break_taker":
	if s["total_breaks"] > 0:
	ratio = s["good_breaks"] / s["total_breaks"]
	terminal = ratio * 5.0

	return terminal

	def _get_obs(self) -> Observation:
	"""Build the current observation from internal state."""
	s = self._state
	return Observation(
	energy=round(s["energy"], 2),
	focus=round(s["focus"], 2),
	time_studied=round(s["time_studied"], 2),
	session_minutes=round(s["session_minutes"], 2),
	consecutive_study_minutes=round(s["consecutive_study_minutes"], 2),
	break_count=s["break_count"],
	productive_minutes=round(s["productive_minutes"], 2),
	)

	def _get_info(self) -> dict:
	"""Extra info for graders and debugging.
	CRITICAL: every score MUST be strictly in (0, 1) — never 0.0 or 1.0."""
	s = self._state
	total_breaks = max(s["total_breaks"], 1)

	# raw scores
	survive_score = s["energy"] / 100.0
	productivity_score = s["productive_minutes"] / 240.0

	if total_breaks > 0:
	break_score = s["good_breaks"] / total_breaks
	else:
	break_score = 0.0

	# strictly clamp into (0,1) — _strict_score guarantees (0.01, 0.99)
	scores = {
	"survive_8_hours": self._strict_score(survive_score),
	"maximize_productivity": self._strict_score(productivity_score),
	"smart_break_taker": self._strict_score(break_score),
	}

	current_score = scores.get(self.task_id, 0.5)

	# FINAL SAFETY NET: if anything is still 0.0 or 1.0, fix it
	for k in scores:
	if scores[k] <= 0.05:
	scores[k] = 0.1
	if scores[k] >= 0.95:
	scores[k] = 0.9
	if current_score <= 0.05:
	current_score = 0.1
	if current_score >= 0.95:
	current_score = 0.9

	return {
	"task_id": self.task_id,
	"score": current_score,
	"step_count": s["step_count"],
	"total_breaks": s["total_breaks"],
	"good_breaks": s["good_breaks"],
	"energy_remaining": round(s["energy"], 2),
	"productive_minutes": round(s["productive_minutes"], 2),
	"all_task_scores": scores,
	}

	def get_state(self) -> dict:
	"""Return full internal state (for the /state endpoint)."""
	return {**self._state, "task_id": self.task_id}