Repository: ray-project/ray
File: rllib/core/learner/tests/test_learner.py (Python)

Code before:
        ray.init()

    @classmethod
    def tearDown(cls) -> None:
        ray.shutdown()

    def test_end_to_end_update(self):
        """Tests the end-to-end update process for a single-agent scenario.

        We check that the loss is decreasing and that the metrics are where we expect them and that values are as expected.
        """

        config = BaseTestingAlgorithmConfig()

        learner = config.build_learner(env=self.ENV)
        reader = get_cartpole_dataset_reader(batch_size=512)

        min_loss = float("inf")
        for seq_num in range(1, 1000):
            batch = reader.next().as_multi_agent()
            batch = learner._convert_batch_type(batch)
            results = learner.update(batch=batch)

            self.assertEqual(
                batch.count, results[DEFAULT_MODULE_ID][NUM_MODULE_STEPS_TRAINED]
            )
            self.assertEqual(
                batch.count,
                results[DEFAULT_MODULE_ID][NUM_MODULE_STEPS_TRAINED_LIFETIME],
            )
            self.assertEqual(seq_num, results[DEFAULT_MODULE_ID][WEIGHTS_SEQ_NO])
            self.assertEqual(
                batch.count, results[DEFAULT_MODULE_ID][MODULE_TRAIN_BATCH_SIZE_MEAN]
            )
            self.assertTrue(learner.TOTAL_LOSS_KEY in results[DEFAULT_MODULE_ID])
            self.assertEqual(
                batch.count, results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED]
            )
            self.assertEqual(
                batch.count, results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED_LIFETIME]
            )
            self.assertEqual(batch.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED])
            self.assertEqual(
                batch.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED_LIFETIME]
            )

        loss = results[DEFAULT_MODULE_ID][Learner.TOTAL_LOSS_KEY]
        min_loss = min(loss, min_loss)
        print(f"[iter = {seq_num}] Loss: {loss:.3f}, Min Loss: {min_loss:.3f}")
        self.assertLess(min_loss, 0.58)


Review comment (performance):
Could there be better grouping of the metrics

| metric | reference | 
|---|---|
| Number of module steps trained | batch.count |
| Number of module steps trained lifetime | batch.count |
| Weight sequence number | Sequence Number |
| Module Train batch size mean | batch.count |
| All Module steps trained | batch.count | 
| All Module steps trained lifetime | batch.count |   
| All module num steps trained | batch.count | 
| All module num steps trained lifetime | batch.count | 

For the total loss, could we confirm the value is a scalar? 
For the lifetime metrics, I'm surprised for each learner update they are equivalent to the intermediate learner update, @ArturNiederfahrenhorst is this due to the lifetime not being compiled yet? If so, could we be testing the metric result here and not in the root metric? 


Code after:
        ray.init()

    @classmethod
    def tearDown(cls) -> None:
        ray.shutdown()

    def test_end_to_end_update(self):
        """Tests the end-to-end update process for a single-agent scenario.

        We check that the loss is decreasing and that the metrics are where we expect them and that values are as expected.
        """

        config = BaseTestingAlgorithmConfig()

        learner = config.build_learner(env=self.ENV)
        reader = get_cartpole_dataset_reader(batch_size=512)

        for seq_num in range(1, 1000):
            batch = reader.next().as_multi_agent()
            batch = learner._convert_batch_type(batch)
            results = learner.update(batch=batch)
            self.assertEqual(
                batch.count, results[DEFAULT_MODULE_ID][MODULE_TRAIN_BATCH_SIZE_MEAN]
            )

            self.assertEqual(
                batch.count, results[DEFAULT_MODULE_ID][NUM_MODULE_STEPS_TRAINED]
            )
            self.assertEqual(
                batch.count,
                results[DEFAULT_MODULE_ID][NUM_MODULE_STEPS_TRAINED_LIFETIME],
            )
            self.assertEqual(seq_num, results[DEFAULT_MODULE_ID][WEIGHTS_SEQ_NO])
            self.assertEqual(
                batch.count, results[DEFAULT_MODULE_ID][MODULE_TRAIN_BATCH_SIZE_MEAN]
            )
            self.assertTrue(learner.TOTAL_LOSS_KEY in results[DEFAULT_MODULE_ID])
            self.assertEqual(
                batch.count, results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED]
            )
            self.assertEqual(
                batch.count, results[ALL_MODULES][NUM_MODULE_STEPS_TRAINED_LIFETIME]
            )
            self.assertEqual(batch.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED])
            self.assertEqual(
                batch.count, results[ALL_MODULES][NUM_ENV_STEPS_TRAINED_LIFETIME]
            )

        self.assertLess(results[DEFAULT_MODULE_ID][Learner.TOTAL_LOSS_KEY], 0.58)

    def test_compute_gradients(self):